diff --git a/ocelot/include/ocelot/executive/CooperativeThreadArray.h b/ocelot/include/ocelot/executive/CooperativeThreadArray.h index 84e51a00..a6452f0a 100644 --- a/ocelot/include/ocelot/executive/CooperativeThreadArray.h +++ b/ocelot/include/ocelot/executive/CooperativeThreadArray.h @@ -475,6 +475,7 @@ namespace executive { void eval_Ex2(CTAContext &context, const ir::PTXInstruction &instr); void eval_Exit(CTAContext &context, const ir::PTXInstruction &instr); void eval_Fma(CTAContext &context, const ir::PTXInstruction &instr); + void eval_Mma(CTAContext &context, const ir::PTXInstruction &instr); void eval_Isspacep(CTAContext &context, const ir::PTXInstruction &instr); void eval_Ld(CTAContext &context, const ir::PTXInstruction &instr); @@ -562,4 +563,3 @@ namespace executive { } #endif - diff --git a/ocelot/include/ocelot/ir/PTXInstruction.h b/ocelot/include/ocelot/ir/PTXInstruction.h index efa813a8..180d612d 100644 --- a/ocelot/include/ocelot/ir/PTXInstruction.h +++ b/ocelot/include/ocelot/ir/PTXInstruction.h @@ -54,6 +54,7 @@ namespace ir { Mad24, Mad, MadC, + Mma, Max, Membar, Min, @@ -143,6 +144,12 @@ namespace ir { None = 0, CC = 1 }; + + enum MmaShape { + MmaM16N8K8, + MmaM16N8K16, + MmaShape_Invalid + }; enum Volatility { Nonvolatile = 0, @@ -426,6 +433,9 @@ namespace ir { /*! indicates data type of instruction */ PTXOperand::DataType type; + /*! Shape for MMA instructions */ + MmaShape mmaShape; + /*! Flag containing one or more floating-point modifiers */ unsigned int modifier; @@ -602,4 +612,3 @@ namespace ir { } #endif - diff --git a/ocelot/include/ocelot/ir/PTXOperand.h b/ocelot/include/ocelot/ir/PTXOperand.h index 853b719c..db9e24b1 100644 --- a/ocelot/include/ocelot/ir/PTXOperand.h +++ b/ocelot/include/ocelot/ir/PTXOperand.h @@ -25,6 +25,7 @@ namespace ir { typedef int32_t PTXS32; typedef int64_t PTXS64; + typedef _Float16 PTXF16; typedef float PTXF32; typedef double PTXF64; @@ -62,6 +63,8 @@ namespace ir { u64, f16, f32, + bf16, + tf32, f64, b8, b16, @@ -285,4 +288,3 @@ namespace std { } #endif - diff --git a/ocelot/include/ocelot/parser/PTXParser.h b/ocelot/include/ocelot/parser/PTXParser.h index 1527b55c..c713e449 100644 --- a/ocelot/include/ocelot/parser/PTXParser.h +++ b/ocelot/include/ocelot/parser/PTXParser.h @@ -137,6 +137,7 @@ namespace parser private: void _setImmediateTypes(); + void _setMovVectorImmediateTypes(); std::string _nameInContext( const std::string& name ); OperandWrapper* _getOperand( const std::string& name ); @@ -236,6 +237,7 @@ namespace parser void constantOperand( double value ); void indexedOperand( const std::string& name, YYLTYPE& location, long long int value ); + void vectorOperand( unsigned int elements ); void addressableOperand( const std::string& name, long long int value, YYLTYPE& location, bool invert ); @@ -267,6 +269,7 @@ namespace parser void instruction(); void instruction( const std::string& opcode, int dataType ); void instruction( const std::string& opcode ); + void mma( int shape, int accumulatorType, int aType, int bType, int cType ); void tex( int dataType ); void tld4( int dataType ); void callPrototypeName( const std::string& identifier ); @@ -366,4 +369,3 @@ namespace parser } #endif - diff --git a/ocelot/src/executive/CooperativeThreadArray.cpp b/ocelot/src/executive/CooperativeThreadArray.cpp index 81acd6dd..aa057b03 100644 --- a/ocelot/src/executive/CooperativeThreadArray.cpp +++ b/ocelot/src/executive/CooperativeThreadArray.cpp @@ -372,6 +372,59 @@ static ir::PTXF32 ftz(int modifier, ir::PTXF32 f) { return f; } +static ir::PTXU16 ftzF16(int modifier, ir::PTXU16 bits) +{ + const bool subnormal = (bits & 0x7c00u) == 0 && (bits & 0x03ffu) != 0; + if ((modifier & ir::PTXInstruction::ftz) && subnormal) { + return bits & 0x8000u; // preserve sign: +0 or -0 + } + return bits; +} + +static ir::PTXF32 f16ToF32(ir::PTXU16 bits) { + ir::PTXF16 half; + std::memcpy(&half, &bits, sizeof(half)); + return static_cast(half); +} + +static ir::PTXF32 tf32FromF32(ir::PTXF32 value) { + ir::PTXU32 bits = hydrazine::bit_cast(value); + if ((bits & 0x7f800000u) == 0x7f800000u) return value; + const ir::PTXU32 discarded = bits & 0x1fffu; + bits &= ~0x1fffu; + if (discarded > 0x1000u || + (discarded == 0x1000u && (bits & 0x2000u))) { + bits += 0x2000u; + } + return hydrazine::bit_cast(bits); +} + +template< typename Source > +static ir::PTXU16 toF16(Source value, int modifier); + +static ir::PTXF32 bf16ToF32(ir::PTXU16 bits) { + return hydrazine::bit_cast( + static_cast(bits) << 16); +} + +static ir::PTXU16 mmaHalfBits(executive::CooperativeThreadArray& cta, + int threadID, const ir::PTXOperand& operand, unsigned int half) +{ + if (operand.bytes() == 4) { + return static_cast( + (cta.operandAsB32(threadID, operand) >> (half * 16)) & 0xffffu); + } + return cta.operandAsU16(threadID, operand); +} + +static ir::PTXF32 mmaHalf(executive::CooperativeThreadArray& cta, + int threadID, const ir::PTXOperand& operand, unsigned int half, + ir::PTXOperand::DataType type) +{ + ir::PTXU16 bits = mmaHalfBits(cta, threadID, operand, half); + return type == ir::PTXOperand::bf16 ? bf16ToF32(bits) : f16ToF32(bits); +} + void executive::CooperativeThreadArray::trace() { if (traceEvents) { currentEvent.contextStackSize = @@ -509,6 +562,8 @@ void executive::CooperativeThreadArray::execute(int PC) { eval_Exit(context, instr); break; case ir::PTXInstruction::Fma: eval_Fma(context, instr); break; + case ir::PTXInstruction::Mma: + eval_Mma(context, instr); break; case ir::PTXInstruction::Isspacep: eval_Isspacep(context, instr); break; case ir::PTXInstruction::Ld: @@ -1717,7 +1772,19 @@ void executive::CooperativeThreadArray::eval_Abs(CTAContext &context, void executive::CooperativeThreadArray::eval_Add(CTAContext &context, const ir::PTXInstruction &instr) { trace(); - if (instr.type == ir::PTXOperand::f32) { + if (instr.type == ir::PTXOperand::f16) { + for (int threadID = 0; threadID < threadCount; threadID++) { + if (!context.predicated(threadID, instr)) continue; + + ir::PTXF32 a = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.a))); + ir::PTXF32 b = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.b))); + ir::PTXU16 d = toF16(a + b, instr.modifier); + setRegAsB16(threadID, instr.d.reg, ftzF16(instr.modifier, d)); + } + } + else if (instr.type == ir::PTXOperand::f32) { for (int threadID = 0; threadID < threadCount; threadID++) { if (!context.predicated(threadID, instr)) continue; ir::PTXF32 d, a = ftz(instr.modifier, operandAsF32(threadID, instr.a)), @@ -2944,6 +3011,41 @@ static ir::PTXF64 toF64(Int value, int modifier) { return d; } +template< typename Source > +static ir::PTXU16 toF16(Source value, int modifier) { + int mode = hydrazine::fegetround(); + if (modifier & ir::PTXInstruction::rn) { + hydrazine::fesetround(FE_TONEAREST); + } else if (modifier & ir::PTXInstruction::rz) { + hydrazine::fesetround(FE_TOWARDZERO); + } else if (modifier & ir::PTXInstruction::rm) { + hydrazine::fesetround(FE_DOWNWARD); + } else if (modifier & ir::PTXInstruction::rp) { + hydrazine::fesetround(FE_UPWARD); + } + ir::PTXF16 half = value; + hydrazine::fesetround(mode); + ir::PTXU16 bits; + std::memcpy(&bits, &half, sizeof(bits)); + return bits; +} + +static ir::PTXU16 f32ToBF16Rn(ir::PTXF32 value) { + ir::PTXU32 bits = hydrazine::bit_cast(value); + if ((bits & 0x7fffffffU) > 0x7f800000U) { + // NVIDIA canonical NaN + return 0x7fffU; + } + ir::PTXU16 upper = static_cast(bits >> 16); + ir::PTXU16 lower = static_cast(bits & 0xffffU); + if (lower > 0x8000U) { + return upper + 1U; + } else if (lower < 0x8000U) { + return upper; + } + return upper + (upper & 1U); +} + template< typename Float > static Float roundToInt(Float a, int modifier, executive::CTAContext &context, const ir::PTXInstruction &instr) { @@ -2983,6 +3085,13 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::u8: { switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(operandAsB8(threadID, instr.a), + instr.modifier)); + } + break; case ir::PTXOperand::pred: // fall through case ir::PTXOperand::b8: // fall through case ir::PTXOperand::u8: // fall through @@ -3033,6 +3142,13 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::s8: { switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(operandAsS8(threadID, instr.a), + instr.modifier)); + } + break; case ir::PTXOperand::s8: // fall through case ir::PTXOperand::s16: // fall through case ir::PTXOperand::s32: // fall through @@ -3084,6 +3200,13 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::u16: { switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(operandAsB16(threadID, instr.a), + instr.modifier)); + } + break; case ir::PTXOperand::pred: // fall through case ir::PTXOperand::b8: // fall through case ir::PTXOperand::u8: @@ -3151,6 +3274,13 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, { // s16 to one of the following switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(operandAsS16(threadID, instr.a), + instr.modifier)); + } + break; case ir::PTXOperand::s8: { ir::PTXS16 a = operandAsS16(threadID, instr.a); @@ -3217,6 +3347,13 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::u32: { switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(operandAsU32(threadID, instr.a), + instr.modifier)); + } + break; case ir::PTXOperand::pred: // fall through case ir::PTXOperand::b8: // fall through case ir::PTXOperand::u8: @@ -3298,6 +3435,13 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::s32: { switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(operandAsS32(threadID, instr.a), + instr.modifier)); + } + break; case ir::PTXOperand::pred: // fall through case ir::PTXOperand::b8: // fall through case ir::PTXOperand::u8: @@ -3378,6 +3522,13 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::s64: { switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(operandAsS64(threadID, instr.a), + instr.modifier)); + } + break; case ir::PTXOperand::pred: // fall through case ir::PTXOperand::b8: // fall through case ir::PTXOperand::u8: // fall through @@ -3474,6 +3625,13 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::u64: { switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(operandAsU64(threadID, instr.a), + instr.modifier)); + } + break; case ir::PTXOperand::pred: // fall through case ir::PTXOperand::b8: // fall through case ir::PTXOperand::u8: @@ -3567,14 +3725,23 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, } } break; + case ir::PTXOperand::f16: // fall through case ir::PTXOperand::f32: { + ir::PTXF32 a = sourceType == ir::PTXOperand::f16 + ? f16ToF32(operandAsU16(threadID, instr.a)) + : operandAsF32(threadID, instr.a); switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(a, instr.modifier)); + } + break; case ir::PTXOperand::pred: // fall through case ir::PTXOperand::b8: // fall through case ir::PTXOperand::u8: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); if (a != a) a = 0.0f; ir::PTXF32 fd = roundToInt(a, instr.modifier, context, instr); @@ -3594,7 +3761,6 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::b16: // fall through case ir::PTXOperand::u16: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); if (a != a) a = 0.0f; ir::PTXF32 fd = roundToInt(a, instr.modifier, context, instr); @@ -3614,7 +3780,6 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::b32: // fall through case ir::PTXOperand::u32: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); if (a != a) a = 0.0f; ir::PTXF32 fd = roundToInt(a, instr.modifier, context, instr); @@ -3634,7 +3799,6 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::b64: // fall through case ir::PTXOperand::u64: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); if (a != a) a = 0.0f; ir::PTXF32 fd = roundToInt(a, instr.modifier, context, instr); @@ -3653,7 +3817,6 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, break; case ir::PTXOperand::s8: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); if (a != a) a = 0.0f; ir::PTXF32 fd = roundToInt(a, instr.modifier, context, instr); @@ -3672,7 +3835,6 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, break; case ir::PTXOperand::s16: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); if (a != a) a = 0.0f; ir::PTXF32 fd = roundToInt(a, instr.modifier, context, instr); @@ -3691,7 +3853,6 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, break; case ir::PTXOperand::s32: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); if (a != a) a = 0.0f; ir::PTXF32 fd = roundToInt(a, instr.modifier, context, instr); @@ -3710,7 +3871,6 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, break; case ir::PTXOperand::s64: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); if (a != a) a = 0.0f; ir::PTXF32 fd = roundToInt(a, instr.modifier, context, instr); @@ -3729,8 +3889,6 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, break; case ir::PTXOperand::f32: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); - a = roundToInt(a, instr.modifier, context, instr); @@ -3740,11 +3898,44 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, break; case ir::PTXOperand::f64: { - ir::PTXF32 a = operandAsF32(threadID, instr.a); ir::PTXF64 d = toF64(a, instr.modifier); setRegAsF64(threadID, instr.d.reg, d); } break; + case ir::PTXOperand::bf16: + { + if (instr.modifier != ir::PTXInstruction::rn) { + throw RuntimeException( + "only cvt.rn.bf16.f32 is implemented", + context.PC, instr); + } + ir::PTXU16 d = f32ToBF16Rn(a); + setRegAsB16(threadID, instr.d.reg, d); + } + break; + default: + throw RuntimeException("conversion not implemented", + context.PC, instr); + break; + } + } + break; + case ir::PTXOperand::bf16: + { + switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(bf16ToF32(operandAsU16(threadID, + instr.a)), instr.modifier)); + } + break; + case ir::PTXOperand::f32: + { + setRegAsF32(threadID, instr.d.reg, + bf16ToF32(operandAsU16(threadID, instr.a))); + } + break; default: throw RuntimeException("conversion not implemented", context.PC, instr); @@ -3755,6 +3946,13 @@ void executive::CooperativeThreadArray::eval_Cvt(CTAContext &context, case ir::PTXOperand::f64: { switch (instr.type) { + case ir::PTXOperand::f16: + { + setRegAsB16(threadID, instr.d.reg, + toF16(operandAsF64(threadID, instr.a), + instr.modifier)); + } + break; case ir::PTXOperand::pred: // fall through case ir::PTXOperand::b8: // fall through case ir::PTXOperand::u8: @@ -4382,6 +4580,16 @@ void executive::CooperativeThreadArray::eval_Ex2(CTAContext &context, setRegAsF32(threadID, instr.d.reg, d); } } + else if (instr.type == ir::PTXOperand::f16) { + for (int threadID = 0; threadID < threadCount; threadID++) { + if (!context.predicated(threadID, instr)) continue; + + ir::PTXF32 a = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.a))); + setRegAsB16(threadID, instr.d.reg, + toF16(hydrazine::exp2f(a), instr.modifier)); + } + } else { throw RuntimeException("unsupported data type", context.PC, instr); } @@ -4426,11 +4634,164 @@ void executive::CooperativeThreadArray::eval_Fma(CTAContext &context, setRegAsF64(tid, instr.d.reg, d); } } + else if (instr.type == ir::PTXOperand::f16) { + for (int tid = 0; tid < threadCount; tid++) { + if (!context.predicated(tid, instr)) continue; + + ir::PTXF32 a = f16ToF32(ftzF16(instr.modifier, + operandAsU16(tid, instr.a))); + ir::PTXF32 b = f16ToF32(ftzF16(instr.modifier, + operandAsU16(tid, instr.b))); + ir::PTXF32 c = f16ToF32(ftzF16(instr.modifier, + operandAsU16(tid, instr.c))); + setRegAsB16(tid, instr.d.reg, + toF16(std::fma(a, b, c), instr.modifier)); + } + } + else if (instr.type == ir::PTXOperand::bf16) { + for (int tid = 0; tid < threadCount; tid++) { + if (!context.predicated(tid, instr)) continue; + + ir::PTXF32 a = bf16ToF32(operandAsU16(tid, instr.a)); + ir::PTXF32 b = bf16ToF32(operandAsU16(tid, instr.b)); + ir::PTXF32 c = bf16ToF32(operandAsU16(tid, instr.c)); + ir::PTXU16 d = f32ToBF16Rn(std::fma(a, b, c)); + setRegAsB16(tid, instr.d.reg, d); + } + } else { throw RuntimeException("unsupported data type", context.PC, instr); } } +void executive::CooperativeThreadArray::eval_Mma(CTAContext &context, + const ir::PTXInstruction &instr) { + trace(); + + if (threadCount < 32 || threadCount % 32 != 0) { + throw RuntimeException("mma requires complete 32-thread warps", + context.PC, instr); + } + + const ir::PTXOperand::DataType inputType = instr.a.type; + const bool halfAccumulator = instr.type == ir::PTXOperand::f16; + const bool tf32Input = inputType == ir::PTXOperand::tf32; + const bool m16n8k8 = instr.mmaShape == ir::PTXInstruction::MmaM16N8K8; + for (int warpStart = 0; warpStart < threadCount; warpStart += 32) { + int participants = 0; + for (int lane = 0; lane < 32; ++lane) { + if (context.predicated(warpStart + lane, instr)) { + ++participants; + } + } + if (participants == 0) continue; + if (participants != 32) { + throw RuntimeException("mma requires all warp lanes to participate", + context.PC, instr); + } + + ir::PTXF32 A[16][16] = {}; + ir::PTXF32 B[16][8] = {}; + ir::PTXF32 C[16][8] = {}; + + for (int lane = 0; lane < 32; ++lane) { + int threadID = warpStart + lane; + int groupID = lane >> 2; + int threadInGroup = lane & 3; + + if (tf32Input) { + for (int i = 0; i < 4; ++i) { + int row = (i & 1) ? groupID + 8 : groupID; + int col = threadInGroup + (i >= 2 ? 4 : 0); + A[row][col] = tf32FromF32(operandAsF32(threadID, + instr.a.array[i])); + } + + for (int i = 0; i < 2; ++i) { + int row = threadInGroup + (i >= 1 ? 4 : 0); + int col = groupID; + B[row][col] = tf32FromF32(operandAsF32(threadID, + instr.b.array[i])); + } + } + else if (m16n8k8) { + for (int i = 0; i < 4; ++i) { + int row = i < 2 ? groupID : groupID + 8; + int col = threadInGroup * 2 + (i & 1); + A[row][col] = mmaHalf(*this, threadID, + instr.a.array[i / 2], i & 1, inputType); + } + + for (int i = 0; i < 2; ++i) { + int row = threadInGroup * 2 + (i & 1); + B[row][groupID] = mmaHalf(*this, threadID, + instr.b.array[0], i, inputType); + } + } + else { + for (int i = 0; i < 8; ++i) { + int row = (i < 2 || (i >= 4 && i < 6)) ? groupID : groupID + 8; + int col = threadInGroup * 2 + (i & 1) + (i >= 4 ? 8 : 0); + A[row][col] = mmaHalf(*this, threadID, + instr.a.array[i / 2], i & 1, inputType); + } + + for (int i = 0; i < 4; ++i) { + int row = threadInGroup * 2 + (i & 1) + (i >= 2 ? 8 : 0); + int col = groupID; + B[row][col] = mmaHalf(*this, threadID, + instr.b.array[i / 2], i & 1, inputType); + } + } + + for (int i = 0; i < 4; ++i) { + int row = groupID + (i >= 2 ? 8 : 0); + int col = threadInGroup * 2 + (i & 1); + C[row][col] = halfAccumulator + ? mmaHalf(*this, threadID, instr.c.array[i / 2], i & 1, + ir::PTXOperand::f16) + : operandAsF32(threadID, instr.c.array[i]); + } + } + + ir::PTXF32 D[16][8]; + for (int row = 0; row < 16; ++row) { + for (int col = 0; col < 8; ++col) { + D[row][col] = C[row][col]; + for (int k = 0; k < (m16n8k8 ? 8 : 16); ++k) { + D[row][col] = std::fma(A[row][k], B[k][col], D[row][col]); + if (halfAccumulator) { + D[row][col] = f16ToF32(toF16(D[row][col], instr.modifier)); + } + } + } + } + + for (int lane = 0; lane < 32; ++lane) { + int threadID = warpStart + lane; + int groupID = lane >> 2; + int threadInGroup = lane & 3; + if (halfAccumulator) { + ir::PTXU32 packed[2] = {0, 0}; + for (int i = 0; i < 4; ++i) { + int row = groupID + (i >= 2 ? 8 : 0); + int col = threadInGroup * 2 + (i & 1); + packed[i / 2] |= static_cast( + toF16(D[row][col], instr.modifier)) << (16 * (i & 1)); + } + setRegAsB32(threadID, instr.d.array[0].reg, packed[0]); + setRegAsB32(threadID, instr.d.array[1].reg, packed[1]); + } + else for (int i = 0; i < 4; ++i) { + int row = groupID + (i >= 2 ? 8 : 0); + int col = threadInGroup * 2 + (i & 1); + setRegAsF32(threadID, instr.d.array[i].reg, + D[row][col]); + } + } + } +} + /*! @@ -5362,6 +5723,33 @@ void executive::CooperativeThreadArray::eval_Max(CTAContext &context, setRegAsF64(threadID, instr.d.reg, d); } } + else if (instr.type == ir::PTXOperand::f16) { + for (int threadID = 0; threadID < threadCount; threadID++) { + if (!context.predicated(threadID, instr)) continue; + + ir::PTXF32 a = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.a))); + ir::PTXF32 b = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.b))); + ir::PTXF32 d; + + if(hydrazine::isnan(a)) + { + d = b; + } + else if(hydrazine::isnan(b)) + { + d = a; + } + else + { + d = ftz(instr.modifier, a > b ? a : b); + } + + setRegAsB16(threadID, instr.d.reg, + toF16(d, instr.modifier)); + } + } else if (instr.type == ir::PTXOperand::s16) { for (int threadID = 0; threadID < threadCount; threadID++) { if (!context.predicated(threadID, instr)) continue; @@ -5813,6 +6201,12 @@ void executive::CooperativeThreadArray::eval_Mov_imm(CTAContext &context, setRegAsU16(threadID, instr.d.reg, a); } break; + case ir::PTXOperand::f16: + { + ir::PTXU16 a = operandAsU16(threadID, instr.a); + setRegAsB16(threadID, instr.d.reg, a); + } + break; case ir::PTXOperand::u32: case ir::PTXOperand::s32: case ir::PTXOperand::b32: @@ -5950,7 +6344,19 @@ void executive::CooperativeThreadArray::eval_Mul24(CTAContext &context, const ir */ void executive::CooperativeThreadArray::eval_Mul(CTAContext &context, const ir::PTXInstruction &instr) { trace(); - if (instr.type == ir::PTXOperand::f32) { + if (instr.type == ir::PTXOperand::f16) { + for (int threadID = 0; threadID < threadCount; threadID++) { + if (!context.predicated(threadID, instr)) continue; + + ir::PTXF32 a = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.a))); + ir::PTXF32 b = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.b))); + ir::PTXU16 d = toF16(a * b, instr.modifier); + setRegAsB16(threadID, instr.d.reg, ftzF16(instr.modifier, d)); + } + } + else if (instr.type == ir::PTXOperand::f32) { for (int threadID = 0; threadID < threadCount; threadID++) { if (!context.predicated(threadID, instr)) continue; @@ -6113,7 +6519,27 @@ void executive::CooperativeThreadArray::eval_Mul(CTAContext &context, const ir:: */ void executive::CooperativeThreadArray::eval_Neg(CTAContext &context, const ir::PTXInstruction &instr) { trace(); - if (instr.type == ir::PTXOperand::f32) { + if (instr.type == ir::PTXOperand::bf16) { + for (int threadID = 0; threadID < threadCount; threadID++) { + if (!context.predicated(threadID, instr)) continue; + + ir::PTXF32 a = ftz(instr.modifier, + bf16ToF32(operandAsU16(threadID, instr.a))); + ir::PTXU16 d = f32ToBF16Rn(-a); + setRegAsB16(threadID, instr.d.reg, d); + } + } + else if (instr.type == ir::PTXOperand::f16) { + for (int threadID = 0; threadID < threadCount; threadID++) { + if (!context.predicated(threadID, instr)) continue; + + ir::PTXF32 a = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.a))); + ir::PTXU16 d = toF16(-a, instr.modifier); + setRegAsB16(threadID, instr.d.reg, ftzF16(instr.modifier, d)); + } + } + else if (instr.type == ir::PTXOperand::f32) { for (int threadID = 0; threadID < threadCount; threadID++) { if (!context.predicated(threadID, instr)) continue; @@ -7132,112 +7558,44 @@ void executive::CooperativeThreadArray::eval_SetP(CTAContext &context, } break; - // single-precision float + // floating-point types [widened to double after type-specific FTZ] + case ir::PTXOperand::f16: case ir::PTXOperand::f32: + case ir::PTXOperand::f64: { for (int threadID = 0; threadID < threadCount; threadID++) { if (!context.predicated(threadID, instr)) continue; - ir::PTXF32 a = ftz(instr.modifier, operandAsF32(threadID, instr.a)), - b = ftz(instr.modifier, operandAsF32(threadID, instr.b)); - bool c = true; // read operator somehow - bool t = false; - - if (instr.c.addressMode == ir::PTXOperand::Register) { - c = operandAsPredicate(threadID, instr.c); - } - - // any branch predictor worth its salt will get this wrong twice or less - switch (instr.comparisonOperator) { - case ir::PTXInstruction::Equ: - case ir::PTXInstruction::Eq: - t = (a == b); - break; - case ir::PTXInstruction::Neu: - case ir::PTXInstruction::Ne: - t = (a != b); - break; - - case ir::PTXInstruction::Ltu: - case ir::PTXInstruction::Lo: // fall through - case ir::PTXInstruction::Lt: - t = (a < b); - break; - - case ir::PTXInstruction::Leu: - case ir::PTXInstruction::Ls: // fall through - case ir::PTXInstruction::Le: - t = (a <= b); - break; - - case ir::PTXInstruction::Gtu: - case ir::PTXInstruction::Hi: // fall through - case ir::PTXInstruction::Gt: - t = (a > b); - break; - - case ir::PTXInstruction::Geu: - case ir::PTXInstruction::Hs: // fall through - case ir::PTXInstruction::Ge: - t = (a >= b); - break; - - case ir::PTXInstruction::Num: - t = !hydrazine::isnan(a) && !hydrazine::isnan(b); - break; - case ir::PTXInstruction::Nan: - t = hydrazine::isnan(a) || hydrazine::isnan(b); - break; - - default: - throw RuntimeException("invalid comparison operator " - "for unsigned int type", context.PC, instr); - } + ir::PTXF64 a; + ir::PTXF64 b; - // now apply the bool op - bool p = false, q = false; - switch (instr.booleanOperator) { - case ir::PTXInstruction::BoolAnd: - p = (t && c); - q = (!t && c); - break; - case ir::PTXInstruction::BoolOr: - p = (t || c); - q = (!t || c); + switch (instr.type) { + case ir::PTXOperand::f16: + a = f16ToF32(ftzF16(instr.modifier, operandAsU16(threadID, instr.a))); + b = f16ToF32(ftzF16(instr.modifier, operandAsU16(threadID, instr.b))); break; - case ir::PTXInstruction::BoolXor: - p = (t && !c) || (!t && c); - q = (!t && !c) || (t && c); + case ir::PTXOperand::f32: + a = ftz(instr.modifier, operandAsF32(threadID, instr.a)); + b = ftz(instr.modifier, operandAsF32(threadID, instr.b)); break; default: - p = t; - q = !t; + a = operandAsF64(threadID, instr.a); + b = operandAsF64(threadID, instr.b); break; } - reportE(REPORT_SETP, " " << instr.a.identifier << " = " << a - << ", " << instr.b.identifier << " = " << b - << " condition = " << t << ", input = " << c << " " - << instr.d.identifier << " = " << p << ", q = " << q ); - - setRegAsPredicate(threadID, instr.d.reg, p); - if (instr.pq.addressMode != ir::PTXOperand::Invalid) { - setRegAsPredicate(threadID, instr.pq.reg, q); - } - } - } - break; + bool c = true; // read operator somehow + bool t = false; - // double-precision float - case ir::PTXOperand::f64: - { - for (int threadID = 0; threadID < threadCount; threadID++) { - if (!context.predicated(threadID, instr)) continue; + const bool hasNaN = hydrazine::isnan(a) || hydrazine::isnan(b); - ir::PTXF64 a = operandAsF64(threadID, instr.a), - b = operandAsF64(threadID, instr.b); - bool c = true; - bool t = false; + const bool unorderedOp = + instr.comparisonOperator == ir::PTXInstruction::Equ || + instr.comparisonOperator == ir::PTXInstruction::Neu || + instr.comparisonOperator == ir::PTXInstruction::Ltu || + instr.comparisonOperator == ir::PTXInstruction::Leu || + instr.comparisonOperator == ir::PTXInstruction::Gtu || + instr.comparisonOperator == ir::PTXInstruction::Geu; if (instr.c.addressMode == ir::PTXOperand::Register) { c = operandAsPredicate(threadID, instr.c); @@ -7245,39 +7603,37 @@ void executive::CooperativeThreadArray::eval_SetP(CTAContext &context, // any branch predictor worth its salt will get this wrong twice or less switch (instr.comparisonOperator) { - case ir::PTXInstruction::Equ: case ir::PTXInstruction::Eq: - t = (a == b); + t = hasNaN ? unorderedOp : (a == b); break; - case ir::PTXInstruction::Neu: case ir::PTXInstruction::Ne: - t = (a != b); + t = hasNaN ? unorderedOp : (a != b); break; case ir::PTXInstruction::Ltu: case ir::PTXInstruction::Lo: // fall through case ir::PTXInstruction::Lt: - t = (a < b); + t = hasNaN ? unorderedOp : (a < b); break; case ir::PTXInstruction::Leu: case ir::PTXInstruction::Ls: // fall through case ir::PTXInstruction::Le: - t = (a <= b); + t = hasNaN ? unorderedOp : (a <= b); break; case ir::PTXInstruction::Gtu: case ir::PTXInstruction::Hi: // fall through case ir::PTXInstruction::Gt: - t = (a > b); + t = hasNaN ? unorderedOp : (a > b); break; case ir::PTXInstruction::Geu: case ir::PTXInstruction::Hs: // fall through case ir::PTXInstruction::Ge: - t = (a >= b); + t = hasNaN ? unorderedOp : (a >= b); break; case ir::PTXInstruction::Num: @@ -7292,7 +7648,6 @@ void executive::CooperativeThreadArray::eval_SetP(CTAContext &context, "for unsigned int type", context.PC, instr); } - // now apply the bool op bool p = false, q = false; switch (instr.booleanOperator) { @@ -7550,14 +7905,23 @@ void executive::CooperativeThreadArray::eval_Set(CTAContext &context, } break; - // single-precision float + // floating-point types, with f16 widened before comparison + case ir::PTXOperand::f16: case ir::PTXOperand::f32: { for (int threadID = 0; threadID < threadCount; threadID++) { if (!context.predicated(threadID, instr)) continue; - ir::PTXF32 a = ftz(instr.modifier, operandAsF32(threadID, instr.a)), + ir::PTXF32 a, b; + if (instr.a.type == ir::PTXOperand::f16) { + a = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.a))); + b = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.b))); + } else { + a = ftz(instr.modifier, operandAsF32(threadID, instr.a)); b = ftz(instr.modifier, operandAsF32(threadID, instr.b)); + } bool c = true; // read operator somehow bool t = false; @@ -7642,12 +8006,15 @@ void executive::CooperativeThreadArray::eval_Set(CTAContext &context, break; } - switch (instr.type) { - case ir::PTXOperand::s32: - case ir::PTXOperand::u32: - setRegAsU32(threadID, instr.d.reg, (t ? 0xFFFFFFFF : 0x00)); - break; - case ir::PTXOperand::f32: + switch (instr.type) { + case ir::PTXOperand::s32: + case ir::PTXOperand::u32: + setRegAsU32(threadID, instr.d.reg, (t ? 0xFFFFFFFF : 0x00)); + break; + case ir::PTXOperand::f16: + setRegAsB16(threadID, instr.d.reg, t ? 0x3c00 : 0x0000); + break; + case ir::PTXOperand::f32: setRegAsF32(threadID, instr.d.reg, (t ? 1.0f : 0.0f)); break; default: @@ -8750,7 +9117,19 @@ void executive::CooperativeThreadArray::eval_St(CTAContext &context, void executive::CooperativeThreadArray::eval_Sub(CTAContext &context, const ir::PTXInstruction &instr) { trace(); - if (instr.type == ir::PTXOperand::f32) { + if (instr.type == ir::PTXOperand::f16) { + for (int threadID = 0; threadID < threadCount; threadID++) { + if (!context.predicated(threadID, instr)) continue; + + ir::PTXF32 a = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.a))); + ir::PTXF32 b = f16ToF32(ftzF16(instr.modifier, + operandAsU16(threadID, instr.b))); + ir::PTXU16 d = toF16(a - b, instr.modifier); + setRegAsB16(threadID, instr.d.reg, ftzF16(instr.modifier, d)); + } + } + else if (instr.type == ir::PTXOperand::f32) { for (int threadID = 0; threadID < threadCount; threadID++) { if (!context.predicated(threadID, instr)) continue; diff --git a/ocelot/src/executive/test/TestInstructions.cpp b/ocelot/src/executive/test/TestInstructions.cpp index ce779bb8..085271c6 100644 --- a/ocelot/src/executive/test/TestInstructions.cpp +++ b/ocelot/src/executive/test/TestInstructions.cpp @@ -19,6 +19,7 @@ #include #include +#include using namespace std; using namespace ir; @@ -42,7 +43,7 @@ class TestInstructions: public Test { status << "Test output:\n"; - threadCount = 16; + threadCount = 32; const std::string ptx = "TestInstructions_ptx"; @@ -323,6 +324,37 @@ class TestInstructions: public Test { PTXInstruction ins; + // f16 + // + if (result) { + ins.opcode = PTXInstruction::Add; + ins.type = PTXOperand::f16; + ins.modifier = PTXInstruction::rn; + ins.a = reg("r1", PTXOperand::b16, 0); + ins.b = reg("r2", PTXOperand::b16, 1); + ins.d = reg("r3", PTXOperand::b16, 2); + + for (int i = 0; i < threadCount; i++) { + cta->setRegAsU16(i, 0, 0x3e00); // 1.5 + cta->setRegAsU16(i, 1, 0x4000); // 2.0 + cta->setRegAsU16(i, 2, 0); + } + if (!ins.valid().empty()) { + result = false; + status << "add.f16 rejected\n"; + } + else { + cta->eval_Add(cta->getActiveContext(), ins); + for (int i = 0; i < threadCount; i++) { + if (cta->getRegAsU16(i, 2) != 0x4300) { // 3.5 + result = false; + status << "add.f16 incorrect\n"; + break; + } + } + } + } + // u16 // if (result) { @@ -597,6 +629,36 @@ class TestInstructions: public Test { PTXInstruction ins; ins.opcode = PTXInstruction::Sub; + // f16 + // + if (result) { + ins.type = PTXOperand::f16; + ins.modifier = PTXInstruction::rn; + ins.a = reg("r1", PTXOperand::b16, 0); + ins.b = reg("r2", PTXOperand::b16, 1); + ins.d = reg("r3", PTXOperand::b16, 2); + + for (int i = 0; i < threadCount; i++) { + cta->setRegAsU16(i, 0, 0x4000); // 2.0 + cta->setRegAsU16(i, 1, 0x3e00); // 1.5 + cta->setRegAsU16(i, 2, 0); + } + if (!ins.valid().empty()) { + result = false; + status << "sub.f16 rejected\n"; + } + else { + cta->eval_Sub(cta->getActiveContext(), ins); + for (int i = 0; i < threadCount; i++) { + if (cta->getRegAsU16(i, 2) != 0x3800) { // 0.5 + result = false; + status << "sub.f16 incorrect\n"; + break; + } + } + } + } + // u16 // if (result) { @@ -1451,6 +1513,62 @@ class TestInstructions: public Test { PTXInstruction ins; ins.opcode = PTXInstruction::Neg; + // bf16 + // + if (result) { + ins.type = PTXOperand::bf16; + ins.modifier = PTXInstruction::rn; + ins.a = reg("r1", PTXOperand::b16, 0); + ins.d = reg("r3", PTXOperand::b16, 2); + + for (int i = 0; i < threadCount; i++) { + cta->setRegAsU16(i, 0, 0x3fc0); // 1.5 + cta->setRegAsU16(i, 2, 0); + } + if (!ins.valid().empty()) { + result = false; + status << "neg.bf16 rejected\n"; + } + else { + cta->eval_Neg(cta->getActiveContext(), ins); + for (int i = 0; i < threadCount; i++) { + if (cta->getRegAsU16(i, 2) != 0xbfc0) { // -1.5 + result = false; + status << "neg.bf16 incorrect\n"; + break; + } + } + } + } + + // f16 + // + if (result) { + ins.type = PTXOperand::f16; + ins.modifier = PTXInstruction::rn; + ins.a = reg("r1", PTXOperand::b16, 0); + ins.d = reg("r3", PTXOperand::b16, 2); + + for (int i = 0; i < threadCount; i++) { + cta->setRegAsU16(i, 0, 0x3e00); // 1.5 + cta->setRegAsU16(i, 2, 0); + } + if (!ins.valid().empty()) { + result = false; + status << "neg.f16 rejected\n"; + } + else { + cta->eval_Neg(cta->getActiveContext(), ins); + for (int i = 0; i < threadCount; i++) { + if (cta->getRegAsU16(i, 2) != 0xbe00) { // -1.5 + result = false; + status << "neg.f16 incorrect\n"; + break; + } + } + } + } + // s16 // if (result) { @@ -2150,6 +2268,39 @@ class TestInstructions: public Test { PTXInstruction ins; ins.opcode = PTXInstruction::Mul; + // f16 + // + if (result) { + ins.type = PTXOperand::f16; + ins.modifier = 0; + ins.a = reg("r1", PTXOperand::b16, 0); + ins.b = reg("r2", PTXOperand::b16, 1); + ins.c = reg("r3", PTXOperand::b16, 2); + ins.d = reg("r4", PTXOperand::b16, 3); + + for (int i = 0; i < threadCount; i++) { + cta->setRegAsU16(i, 0, 0x3e00); // 1.5 + cta->setRegAsU16(i, 1, 0x4000); // 2.0 + cta->setRegAsU16(i, 2, 0); + cta->setRegAsU16(i, 3, 0); + } + std::string error = ins.valid(); + if (!error.empty()) { + result = false; + status << "mul.f16 rejected: " << error << "\n"; + } + else { + cta->eval_Mul(cta->getActiveContext(), ins); + for (int i = 0; i < threadCount; i++) { + if (cta->getRegAsU16(i, 3) != 0x4200) { // 3.0 + result = false; + status << "mul.f16 incorrect\n"; + break; + } + } + } + } + // u16 // if (result) { @@ -2686,6 +2837,195 @@ class TestInstructions: public Test { return result; } + bool test_Bf16Fma() { + PTXInstruction ins; + ins.opcode = PTXInstruction::Fma; + ins.type = PTXOperand::bf16; + ins.modifier = PTXInstruction::rn; + ins.a = reg("r1", PTXOperand::b16, 0); + ins.b = reg("r2", PTXOperand::b16, 1); + ins.c = reg("r4", PTXOperand::b16, 3); + ins.d = reg("r3", PTXOperand::b16, 2); + + // 1.0 * 1.0078125 + 0.00390625 is halfway between + // 0x3f81 and 0x3f82, so round to the even result 0x3f82. + for (int i = 0; i < threadCount; i++) { + cta->setRegAsU16(i, 0, 0x3f80); + cta->setRegAsU16(i, 1, 0x3f81); + cta->setRegAsU16(i, 3, 0x3b80); + cta->setRegAsU16(i, 2, 0); + } + + cta->eval_Fma(cta->getActiveContext(), ins); + for (int i = 0; i < threadCount; i++) { + if (cta->getRegAsU16(i, 2) != 0x3f82) { + status << "fma.rn.bf16 incorrect [" << i << "]\n"; + return false; + } + } + + return true; + } + + bool test_F16Fma() { + PTXInstruction ins; + ins.opcode = PTXInstruction::Fma; + ins.type = PTXOperand::f16; + ins.modifier = PTXInstruction::rn; + ins.a = reg("a", PTXOperand::b16, 0); + ins.b = reg("b", PTXOperand::b16, 1); + ins.c = reg("c", PTXOperand::b16, 3); + ins.d = reg("d", PTXOperand::b16, 2); + + // 1.0 * (1.0 + 2^-10) - 2^-11 is halfway between + // 1.0 and the next half value, so round to the even result 1.0. + for (int i = 0; i < threadCount; i++) { + cta->setRegAsU16(i, 0, 0x3c00); + cta->setRegAsU16(i, 1, 0x3c01); + cta->setRegAsU16(i, 3, 0x9000); + cta->setRegAsU16(i, 2, 0); + } + + cta->eval_Fma(cta->getActiveContext(), ins); + for (int i = 0; i < threadCount; i++) { + if (cta->getRegAsU16(i, 2) != 0x3c00) { + status << "fma.rn.f16 incorrect [" << i << "]\n"; + return false; + } + } + + return true; + } + + bool test_Mma() { + PTXInstruction ins; + ins.opcode = PTXInstruction::Mma; + ins.type = PTXOperand::f32; + ins.modifier = PTXInstruction::rn; + + auto vector = [this](PTXOperand::DataType type, + PTXOperand::DataType elementType, PTXOperand::Vec vec, + int firstRegister, int count) { + PTXOperand operand; + operand.addressMode = PTXOperand::Register; + operand.type = type; + operand.vec = vec; + for(int i = 0; i < count; ++i) { + operand.array.push_back(reg("mma", elementType, + (PTXOperand::RegisterType)(firstRegister + i))); + } + return operand; + }; + + ins.d = vector(PTXOperand::f32, PTXOperand::f32, + PTXOperand::v4, 0, 4); + ins.c = vector(PTXOperand::f32, PTXOperand::f32, + PTXOperand::v4, 6, 4); + ins.a = vector(PTXOperand::f16, PTXOperand::b32, + PTXOperand::v4, 0, 4); + ins.b = vector(PTXOperand::f16, PTXOperand::b32, + PTXOperand::v2, 4, 2); + + const PTXU16 f16Values[17] = { + 0x0000, 0x3c00, 0x4000, 0x4200, 0x4400, 0x4500, + 0x4600, 0x4700, 0x4800, 0x4880, 0x4900, 0x4980, + 0x4a00, 0x4a80, 0x4b00, 0x4b80, 0x4c00 + }; + cta->reset(); + for(int thread = 0; thread < threadCount; ++thread) { + const int lane = thread & 31; + const int groupID = lane >> 2; + const int threadInGroup = lane & 3; + for(int regIndex = 0; regIndex < 4; ++regIndex) { + PTXU32 packed = 0; + for(int half = 0; half < 2; ++half) { + const int i = regIndex * 2 + half; + const int row = (i < 2 || (i >= 4 && i < 6)) + ? groupID : groupID + 8; + const int value = row + 1; + packed |= (PTXU32)f16Values[value] << (16 * half); + } + cta->setRegAsB32(thread, regIndex, packed); + } + for(int regIndex = 0; regIndex < 2; ++regIndex) { + PTXU32 packed = 0; + for(int half = 0; half < 2; ++half) { + const int col = groupID; + packed |= (PTXU32)f16Values[col + 1] << (16 * half); + } + cta->setRegAsB32(thread, 4 + regIndex, packed); + } + for(int i = 0; i < 4; ++i) { + const int row = groupID + (i >= 2 ? 8 : 0); + const int col = threadInGroup * 2 + (i & 1); + cta->setRegAsF32(thread, 6 + i, (PTXF32)(100 * row + col)); + } + } + + cta->eval_Mma(cta->getActiveContext(), ins); + for(int thread = 0; thread < threadCount; ++thread) { + const int lane = thread & 31; + const int groupID = lane >> 2; + const int threadInGroup = lane & 3; + for(int regIndex = 0; regIndex < 4; ++regIndex) { + const int row = groupID + (regIndex >= 2 ? 8 : 0); + const int col = threadInGroup * 2 + (regIndex & 1); + const PTXF32 expected = 16.0f * (row + 1) * (col + 1) + + (PTXF32)(100 * row + col); + if(std::fabs(cta->getRegAsF32(thread, regIndex) - expected) > 0.001f) { + status << "mma.m16n8k16.f16 incorrect [" + << thread << "]\n"; + return false; + } + } + } + + ins.a.type = PTXOperand::bf16; + ins.b.type = PTXOperand::bf16; + const PTXU32 bf16One = 0x3f803f80u; + const PTXU32 bf16Two = 0x40004000u; + cta->reset(); + for(int thread = 0; thread < threadCount; ++thread) { + for(int regIndex = 0; regIndex < 4; ++regIndex) { + cta->setRegAsB32(thread, regIndex, bf16One); + } + for(int regIndex = 4; regIndex < 6; ++regIndex) { + cta->setRegAsB32(thread, regIndex, bf16Two); + } + for(int regIndex = 6; regIndex < 10; ++regIndex) { + cta->setRegAsF32(thread, regIndex, 3.0f); + } + } + + cta->eval_Mma(cta->getActiveContext(), ins); + for(int thread = 0; thread < threadCount; ++thread) { + for(int regIndex = 0; regIndex < 4; ++regIndex) { + if(std::fabs(cta->getRegAsF32(thread, regIndex) - 35.0f) > 0.001f) { + status << "mma.m16n8k16.bf16 incorrect [" + << thread << "]\n"; + return false; + } + } + } + + // MMA is warp-collective: a partially active warp must not execute it. + cta->getActiveContext().active[0] = false; + bool rejectedPartialWarp = false; + try { + cta->eval_Mma(cta->getActiveContext(), ins); + } + catch (RuntimeException &) { + rejectedPartialWarp = true; + } + cta->getActiveContext().active[0] = true; + if (!rejectedPartialWarp) { + status << "mma.m16n8k16 accepted a partial warp\n"; + return false; + } + + return true; + } + bool test_Lg2() { bool result = true; @@ -3784,6 +4124,52 @@ class TestInstructions: public Test { } } + // pack a 16-bit immediate and register into a 32-bit destination + if (result) { + ins.d = reg("f10", PTXOperand::f32, 0); + ins.a = PTXOperand(); + ins.a.addressMode = PTXOperand::Register; + ins.a.type = PTXOperand::s16; + ins.a.vec = PTXOperand::v2; + ins.a.array.push_back(imm_uint("0", PTXOperand::s16, 0)); + ins.a.array.push_back(reg("rs1", PTXOperand::s16, 1)); + ins.type = PTXOperand::b32; + + for (int i = 0; i < threadCount; ++i) { + cta->setRegAsB16(i, 1, 0x3f80); + } + + cta->eval_Mov(cta->getActiveContext(), ins); + + for (int i = 0; i < threadCount; ++i) { + if (cta->getRegAsB32(i, 0) != 0x3f800000) { + result = false; + status << "mov.b32 f10, {0, rs1} failed\n"; + break; + } + } + } + + // pack two 16-bit immediates into a 32-bit destination + if (result) { + ins.a = PTXOperand(); + ins.a.addressMode = PTXOperand::Register; + ins.a.type = PTXOperand::b16; + ins.a.vec = PTXOperand::v2; + ins.a.array.push_back(imm_uint("5", PTXOperand::b16, 5)); + ins.a.array.push_back(imm_uint("3", PTXOperand::b16, 3)); + + cta->eval_Mov(cta->getActiveContext(), ins); + + for (int i = 0; i < threadCount; ++i) { + if (cta->getRegAsB32(i, 0) != 0x00030005) { + result = false; + status << "mov.b32 f10, {5, 3} failed\n"; + break; + } + } + } + // from label return result; @@ -3797,7 +4183,160 @@ class TestInstructions: public Test { cta->reset(); - // + // cvt.rn.bf16.f32 + ins.type = PTXOperand::bf16; + ins.modifier = PTXInstruction::rn; + ins.d = reg("d", PTXOperand::b16, 0); + ins.a = reg("a", PTXOperand::f32, 1); + + const PTXU32 input[] = { + 0x3f800000, // exact + 0x3f807fff, // below halfway + 0x3f808000, // halfway, upper even + 0x3f808001, // above halfway + 0x3f818000, // halfway, upper odd + 0x00000000, // +0 + 0x80000000, // -0 + 0x7f800000, // +infinity + 0x7fc00000 // NaN + }; + const PTXU16 expected[] = { + 0x3f80, + 0x3f80, + 0x3f80, + 0x3f81, + 0x3f82, + 0x0000, + 0x8000, + 0x7f80, + 0x7fff + }; + const int cases = sizeof(input) / sizeof(input[0]); + + for (int i = 0; i < threadCount; ++i) { + cta->setRegAsU32(i, 1, input[i % cases]); + cta->setRegAsU16(i, 0, 0); + } + + cta->eval_Cvt(cta->getActiveContext(), ins); + + for (int i = 0; i < threadCount; ++i) { + PTXU16 got = cta->getRegAsU16(i, 0); + if (got != expected[i % cases]) { + status << "cvt.rn.bf16.f32 failed (thread " << i + << "): expected 0x" << hex << expected[i % cases] + << ", got 0x" << got << dec << "\n"; + result = false; + break; + } + } + + if (result) { + ins.modifier = PTXInstruction::rz; + try { + cta->eval_Cvt(cta->getActiveContext(), ins); + status << "cvt.rz.bf16.f32 should not be implemented\n"; + result = false; + } catch (RuntimeException &) { + // Expected: only .rn is implemented. + } + } + + if (result) { + // cvt.f32.bf16 + ins.type = PTXOperand::f32; + ins.modifier = 0; + ins.d = reg("d", PTXOperand::f32, 0); + ins.a = reg("a", PTXOperand::bf16, 1); + + cta->setRegAsU16(0, 1, 0xc020); + cta->eval_Cvt(cta->getActiveContext(), ins); + + if (cta->getRegAsU32(0, 0) != 0xc0200000) { + status << "cvt.f32.bf16 failed\n"; + result = false; + } + } + + if (result) { + // cvt.rn.f16.f32 + ins.type = PTXOperand::f16; + ins.modifier = PTXInstruction::rn; + ins.d = reg("d", PTXOperand::b16, 0); + ins.a = reg("a", PTXOperand::f32, 1); + + cta->setRegAsF32(0, 1, 2049.0f); + cta->eval_Cvt(cta->getActiveContext(), ins); + + if (cta->getRegAsU16(0, 0) != 0x6800) { + status << "cvt.rn.f16.f32 failed\n"; + result = false; + } + } + + if (result) { + // cvt.f32.f16 + ins.type = PTXOperand::f32; + ins.modifier = 0; + ins.d = reg("d", PTXOperand::f32, 0); + ins.a = reg("a", PTXOperand::f16, 1); + + cta->setRegAsU16(0, 1, 0x3c00); + cta->eval_Cvt(cta->getActiveContext(), ins); + + if (cta->getRegAsU32(0, 0) != 0x3f800000) { + status << "cvt.f32.f16 failed\n"; + result = false; + } + } + + if (result) { + // cvt.rzi.s32.f16 + ins.type = PTXOperand::s32; + ins.modifier = PTXInstruction::rzi; + ins.d = reg("d", PTXOperand::s32, 0); + ins.a = reg("a", PTXOperand::f16, 1); + + cta->setRegAsU16(0, 1, 0x3e00); + cta->eval_Cvt(cta->getActiveContext(), ins); + + if (cta->getRegAsS32(0, 0) != 1) { + status << "cvt.rzi.s32.f16 failed\n"; + result = false; + } + } + + if (result) { + // cvt.rn.f16.s64 + ins.type = PTXOperand::f16; + ins.modifier = PTXInstruction::rn; + ins.d = reg("d", PTXOperand::b16, 0); + ins.a = reg("a", PTXOperand::s64, 1); + + cta->setRegAsS64(0, 1, 2049); + cta->eval_Cvt(cta->getActiveContext(), ins); + + if (cta->getRegAsU16(0, 0) != 0x6800) { + status << "cvt.rn.f16.s64 failed\n"; + result = false; + } + } + + if (result) { + // cvt.rn.f16.f64 + ins.type = PTXOperand::f16; + ins.modifier = PTXInstruction::rn; + ins.d = reg("d", PTXOperand::b16, 0); + ins.a = reg("a", PTXOperand::f64, 1); + + cta->setRegAsF64(0, 1, 1.0); + cta->eval_Cvt(cta->getActiveContext(), ins); + + if (cta->getRegAsU16(0, 0) != 0x3c00) { + status << "cvt.rn.f16.f64 failed\n"; + result = false; + } + } return result; } @@ -3900,6 +4439,35 @@ class TestInstructions: public Test { } } + if (result) { + // set.eq.f16.f16.and + ins = PTXInstruction(); + ins.opcode = PTXInstruction::Set; + ins.type = PTXOperand::f16; + ins.d = reg("d", PTXOperand::b16, 3); + ins.a = reg("a", PTXOperand::f16, 1); + ins.b = reg("b", PTXOperand::f16, 2); + ins.c = reg("c", PTXOperand::pred, 0); + ins.comparisonOperator = PTXInstruction::Eq; + ins.booleanOperator = PTXInstruction::BoolAnd; + + cta->setRegAsU16(0, 1, 0x3c00); // 1.0 + cta->setRegAsU16(0, 2, 0x3c00); // 1.0 + cta->setRegAsPredicate(0, 0, true); + cta->eval_Set(cta->getActiveContext(), ins); + if (cta->getRegAsU16(0, 3) != 0x3c00) { + status << "[set.eq.f16.f16.and test] failed\n"; + result = false; + } + + cta->setRegAsPredicate(0, 0, false); + cta->eval_Set(cta->getActiveContext(), ins); + if (cta->getRegAsU16(0, 3) != 0x0000) { + status << "[set.eq.f16.f16.and false predicate test] failed\n"; + result = false; + } + } + return result; } @@ -4119,6 +4687,65 @@ class TestInstructions: public Test { } } + if (result) { + // Unordered floating-point comparisons are true when an input is NaN. + ins = PTXInstruction(); + ins.opcode = PTXInstruction::SetP; + ins.type = PTXOperand::f64; + ins.d = reg("p", PTXOperand::pred, 3); + ins.pq = reg("q", PTXOperand::pred, 4); + ins.a = reg("a", PTXOperand::f64, 1); + ins.b = reg("b", PTXOperand::f64, 2); + ins.comparisonOperator = PTXInstruction::Equ; + + cta->setRegAsF64(0, 1, + std::numeric_limits::quiet_NaN()); + cta->setRegAsF64(0, 2, 1.0); + cta->eval_SetP(cta->getActiveContext(), ins); + + if (!cta->getRegAsPredicate(0, 3) || + cta->getRegAsPredicate(0, 4)) { + status << "[f64 Equ NaN test] " << ins.toString() + << " failed\n"; + result = false; + } + } + + if (result) { + // Half inputs are widened exactly, with FTZ applied before widening. + ins = PTXInstruction(); + ins.opcode = PTXInstruction::SetP; + ins.type = PTXOperand::f16; + ins.d = reg("p", PTXOperand::pred, 3); + ins.pq = reg("q", PTXOperand::pred, 4); + ins.a = reg("a", PTXOperand::b16, 1); + ins.b = reg("b", PTXOperand::b16, 2); + ins.comparisonOperator = PTXInstruction::Lt; + + cta->setRegAsU16(0, 1, 0x3c00); // 1.0 + cta->setRegAsU16(0, 2, 0x4000); // 2.0 + cta->eval_SetP(cta->getActiveContext(), ins); + const bool normal = cta->getRegAsPredicate(0, 3) && + !cta->getRegAsPredicate(0, 4); + + ins.comparisonOperator = PTXInstruction::Eq; + cta->setRegAsU16(0, 1, 0x0001); // minimum half subnormal + cta->setRegAsU16(0, 2, 0x0000); + cta->eval_SetP(cta->getActiveContext(), ins); + const bool preserved = !cta->getRegAsPredicate(0, 3) && + cta->getRegAsPredicate(0, 4); + + ins.modifier = PTXInstruction::ftz; + cta->eval_SetP(cta->getActiveContext(), ins); + const bool flushed = cta->getRegAsPredicate(0, 3) && + !cta->getRegAsPredicate(0, 4); + + if (!normal || !preserved || !flushed) { + status << "[f16 widening/FTZ test] failed\n"; + result = false; + } + } + return result; } @@ -4443,6 +5070,7 @@ class TestInstructions: public Test { result = (result && test_Mov()); // cvt instruction + result = (result && test_Cvt()); // arithmetic instructions result = (result && test_Abs()); @@ -4470,6 +5098,9 @@ class TestInstructions: public Test { result = (result && test_Cos()); result = (result && test_Sin()); result = (result && test_Ex2()); + result = (result && test_F16Fma()); + result = (result && test_Bf16Fma()); + result = (result && test_Mma()); result = (result && test_Lg2()); result = (result && test_Sqrt()); result = (result && test_Rsqrt()); @@ -4540,4 +5171,3 @@ int main(int argc, char **argv) { return test.passed(); } - diff --git a/ocelot/src/ir/PTXInstruction.cpp b/ocelot/src/ir/PTXInstruction.cpp index 099a0493..fc89c76b 100644 --- a/ocelot/src/ir/PTXInstruction.cpp +++ b/ocelot/src/ir/PTXInstruction.cpp @@ -382,6 +382,7 @@ std::string ir::PTXInstruction::toString( Opcode opcode ) { case Lg2: return "lg2"; break; case Mad24: return "mad24"; break; case Mad: return "mad"; break; + case Mma: return "mma"; break; case MadC: return "madc"; break; case Max: return "max"; break; case Membar: return "membar"; break; @@ -458,6 +459,7 @@ ir::PTXInstruction::PTXInstruction( Opcode op, const PTXOperand& _d, reconvergeInstruction = 0; branchTargetInstruction = 0; vec = PTXOperand::v1; + mmaShape = MmaShape_Invalid; pg.condition = PTXOperand::PT; pg.type = PTXOperand::pred; barrierOperation = BarSync; @@ -504,7 +506,7 @@ std::string ir::PTXInstruction::valid() const { } case Add: { if ( !( type != PTXOperand::s8 && type != PTXOperand::u8 && - type != PTXOperand::b8 && type != PTXOperand::f16 + type != PTXOperand::b8 && type != PTXOperand::pred ) ) { return "invalid instruction type " + PTXOperand::toString( type ); @@ -921,7 +923,7 @@ std::string ir::PTXInstruction::valid() const { break; } case Ex2: { - if( !( type == PTXOperand::f32 ) ) { + if( !( type == PTXOperand::f32 || type == PTXOperand::f16 ) ) { return "invalid instruction type " + PTXOperand::toString( type ); } @@ -945,7 +947,8 @@ std::string ir::PTXInstruction::valid() const { break; } case Fma: { - if (!(type == ir::PTXOperand::f32 || type == ir::PTXOperand::f64)) { + if (!(type == ir::PTXOperand::f16 || type == ir::PTXOperand::f32 + || type == ir::PTXOperand::f64 || type == ir::PTXOperand::bf16)) { return "invalid instruction type " + PTXOperand::toString( type ); } if( !PTXOperand::valid( type, d.type ) ) { @@ -1094,6 +1097,84 @@ std::string ir::PTXInstruction::valid() const { } break; } + case Mma: { + const bool m16n8k8 = mmaShape == MmaM16N8K8; + const bool tf32Input = a.type == PTXOperand::tf32; + const bool halfAccumulator = type == PTXOperand::f16; + if (mmaShape == MmaShape_Invalid) { + return "mma has no shape"; + } + if (!halfAccumulator && type != PTXOperand::f32) { + return "mma requires f16 or f32 accumulators"; + } + if (a.type != PTXOperand::f16 && a.type != PTXOperand::bf16 && + a.type != PTXOperand::tf32) { + return "mma A type must be f16, bf16, or tf32"; + } + if (tf32Input && (!m16n8k8 || type != PTXOperand::f32 || + b.type != PTXOperand::tf32)) { + return "tf32 mma requires m16n8k8 with f32 accumulators"; + } + if (halfAccumulator && a.type != PTXOperand::f16) { + return "f16 mma accumulators require f16 inputs"; + } + if (b.type != a.type) { + return "mma A and B types must match"; + } + if (d.type != type || c.type != type) { + return halfAccumulator ? "mma C and D types must be f16" + : "mma C and D types must be f32"; + } + const PTXOperand::Vec accumulatorVec = halfAccumulator + ? PTXOperand::v2 : PTXOperand::v4; + const unsigned int accumulatorRegisters = halfAccumulator ? 2 : 4; + const bool compactInputFragment = m16n8k8 && !tf32Input; + const PTXOperand::Vec inputAVec = compactInputFragment + ? PTXOperand::v2 : PTXOperand::v4; + const PTXOperand::Vec inputBVec = compactInputFragment + ? PTXOperand::v1 : PTXOperand::v2; + if (d.vec != accumulatorVec || c.vec != accumulatorVec || + a.vec != inputAVec || b.vec != inputBVec) { + return m16n8k8 ? "mma.m16n8k8 has invalid fragment sizes" + : "mma.m16n8k16 has invalid fragment sizes"; + } + if (d.array.size() != accumulatorRegisters || + c.array.size() != accumulatorRegisters || + a.array.size() != (compactInputFragment ? 2u : 4u) || + b.array.size() != (compactInputFragment ? 1u : 2u)) { + return m16n8k8 ? "mma.m16n8k8 has invalid fragment register counts" + : "mma.m16n8k16 has invalid fragment register counts"; + } + for (PTXOperand::Array::const_iterator element = a.array.begin(); + element != a.array.end(); ++element) { + if (element->type != PTXOperand::b32) { + return "mma A fragment registers must be 32-bit packed values"; + } + } + for (PTXOperand::Array::const_iterator element = b.array.begin(); + element != b.array.end(); ++element) { + if (element->type != PTXOperand::b32) { + return "mma B fragment registers must be 32-bit packed values"; + } + } + for (PTXOperand::Array::const_iterator element = c.array.begin(); + element != c.array.end(); ++element) { + if (halfAccumulator ? element->type != PTXOperand::b32 + : !PTXOperand::relaxedValid(PTXOperand::f32, element->type)) { + return halfAccumulator ? "mma C fragment registers must be b32" + : "mma C fragment registers must be f32 or b32"; + } + } + for (PTXOperand::Array::const_iterator element = d.array.begin(); + element != d.array.end(); ++element) { + if (halfAccumulator ? element->type != PTXOperand::b32 + : !PTXOperand::relaxedValid(PTXOperand::f32, element->type)) { + return halfAccumulator ? "mma D fragment registers must be b32" + : "mma D fragment registers must be f32 or b32"; + } + } + break; + } case MadC: { if( !( type == PTXOperand::u32 || type == PTXOperand::s32 ) ) { return "invalid instruction type " @@ -1122,7 +1203,7 @@ std::string ir::PTXInstruction::valid() const { } case Max: { if( !( type != PTXOperand::s8 && type != PTXOperand::u8 && - type != PTXOperand::b8 && type != PTXOperand::f16 + type != PTXOperand::b8 && type != PTXOperand::pred ) ) { return "invalid instruction type " + PTXOperand::toString( type ); @@ -1178,14 +1259,16 @@ std::string ir::PTXInstruction::valid() const { break; } case Mov: { - if ( ( a.type == PTXOperand::f16 ) && + if ( type != PTXOperand::b16 && + a.type == PTXOperand::f16 && + a.array.empty() && a.addressMode != PTXOperand::Address && a.addressMode != PTXOperand::Immediate ) { - return "invalid type for operand A " + return "invalid type for operand A " + PTXOperand::toString( a.type ); } if ( !( d.type != PTXOperand::s8 && d.type != PTXOperand::u8 - && d.type != PTXOperand::b8 && d.type != PTXOperand::f16 ) ) { + && d.type != PTXOperand::b8 ) ) { return "invalid type for operand D " + PTXOperand::toString( d.type ); } @@ -1222,7 +1305,7 @@ std::string ir::PTXInstruction::valid() const { } case Mul: { if( type == PTXOperand::s8 || type == PTXOperand::u8 - || type == PTXOperand::b8 || type == PTXOperand::f16 + || type == PTXOperand::b8 || type == PTXOperand::pred ) { return "invalid instruction type " + PTXOperand::toString( type ); @@ -1264,7 +1347,8 @@ std::string ir::PTXInstruction::valid() const { case Neg: { if( type != PTXOperand::s16 && type != PTXOperand::s32 && type != PTXOperand::s64 && type != PTXOperand::f32 && - type != PTXOperand::f64 ) { + type != PTXOperand::f64 && type != PTXOperand::f16 && + type != PTXOperand::bf16 ) { return "invalid instruction type " + PTXOperand::toString( type ); } @@ -1576,14 +1660,15 @@ std::string ir::PTXInstruction::valid() const { && type != PTXOperand::u32 && type != PTXOperand::u64 && type != PTXOperand::b16 && type != PTXOperand::b32 && type != PTXOperand::b64 && type != PTXOperand::f32 - && type != PTXOperand::f64 ) { + && type != PTXOperand::f64 && type != PTXOperand::f16 ) { return "invalid instruction type " + PTXOperand::toString( type ); } - if( d.type != PTXOperand::s32 && d.type != PTXOperand::f32 - && d.type != PTXOperand::u32 ) { + if( d.type != PTXOperand::s32 && d.type != PTXOperand::f32 + && d.type != PTXOperand::u32 && d.type != PTXOperand::b32 + && d.type != PTXOperand::b16 && d.type != PTXOperand::f16 ) { return "operand D type " + PTXOperand::toString( d.type ) - + " invalid (must be u32, s32, or f32)"; + + " invalid (must be b16, f16, b32, u32, s32, or f32)"; } if( c.type != PTXOperand::pred && c.addressMode != PTXOperand::Invalid ) { @@ -1615,7 +1700,7 @@ std::string ir::PTXInstruction::valid() const { && type != PTXOperand::u32 && type != PTXOperand::u64 && type != PTXOperand::b16 && type != PTXOperand::b32 && type != PTXOperand::b64 && type != PTXOperand::f32 - && type != PTXOperand::f64 ) { + && type != PTXOperand::f64 && type != PTXOperand::f16 ) { return "invalid instruction type " + PTXOperand::toString( type ); } @@ -1862,7 +1947,7 @@ std::string ir::PTXInstruction::valid() const { } case Sub: { if ( !( type != PTXOperand::s8 && type != PTXOperand::u8 && - type != PTXOperand::b8 && type != PTXOperand::f16 + type != PTXOperand::b8 && type != PTXOperand::pred ) ) { return "invalid instruction type " + PTXOperand::toString( type ); @@ -2363,6 +2448,20 @@ std::string ir::PTXInstruction::toString() const { + a.toString() + ", " + b.toString() + ", " + c.toString(); return result; } + case Mma: { + const bool m16n8k8 = mmaShape == MmaM16N8K8; + std::string result = guard() + + "mma.sync.aligned.m" + + (m16n8k8 ? "16n8k8" : "16n8k16") + + ".row.col." + + PTXOperand::toString(type) + "." + + PTXOperand::toString(a.type) + "." + + PTXOperand::toString(b.type) + "." + + PTXOperand::toString(c.type) + " " + + d.toString() + ", " + a.toString() + ", " + + b.toString() + ", " + c.toString(); + return result; + } case MadC: { std::string result = guard() + "madc."; result += modifierString( modifier, carry ); diff --git a/ocelot/src/ir/PTXOperand.cpp b/ocelot/src/ir/PTXOperand.cpp index 7ea6c4c9..fb0d7e99 100644 --- a/ocelot/src/ir/PTXOperand.cpp +++ b/ocelot/src/ir/PTXOperand.cpp @@ -50,6 +50,8 @@ std::string ir::PTXOperand::toString( DataType type ) { case b64: return "b64"; break; case f16: return "f16"; break; case f32: return "f32"; break; + case bf16: return "bf16"; break; + case tf32: return "tf32"; break; case f64: return "f64"; break; case pred: return "pred"; break; default: break; @@ -150,6 +152,7 @@ bool ir::PTXOperand::isFloat( DataType type ) { switch( type ) { case f16: /* fall through */ case f32: /* fall through */ + case bf16:/* fall through */ case f64: result = true; default: break; } @@ -194,10 +197,12 @@ unsigned int ir::PTXOperand::bytes( DataType type ) { case u16: /* fall through */ case f16: /* fall through */ case b16: /* fall through */ + case bf16: /* fall through */ case s16: return 2; break; case u32: /* fall through */ case b32: /* fall through */ case f32: /* fall through */ + case tf32: /* fall through */ case s32: return 4; break; case f64: /* fall through */ case u64: /* fall through */ @@ -235,6 +240,7 @@ bool ir::PTXOperand::valid( DataType destination, DataType source ) { case s16: /* fall through */ case u16: /* fall through */ case f16: /* fall through */ + case bf16: /* fall through */ case b16: return true; break; default: break; } @@ -345,6 +351,10 @@ bool ir::PTXOperand::valid( DataType destination, DataType source ) { } break; } + case bf16: { + return source == b16; + break; + } case pred: { return source == pred; break; @@ -553,6 +563,12 @@ bool ir::PTXOperand::relaxedValid( DataType instructionType, } break; } + case bf16: { + switch( operand ) { + case b16: return true; break; + default: break; + } + } case pred: { return operand == pred; break; @@ -864,5 +880,3 @@ bool ir::PTXOperand::isRegister() const { bool ir::PTXOperand::isVector() const { return isRegister() && vec != v1; } - - diff --git a/ocelot/src/parser/PTXLexer.cpp b/ocelot/src/parser/PTXLexer.cpp index 8d9c3c01..ee62aed3 100644 --- a/ocelot/src/parser/PTXLexer.cpp +++ b/ocelot/src/parser/PTXLexer.cpp @@ -64,6 +64,7 @@ namespace parser CASE(OPCODE_REM) CASE(OPCODE_MUL24) CASE(OPCODE_MAD24) + CASE(OPCODE_MMA) CASE(OPCODE_DIV) CASE(OPCODE_ABS) CASE(OPCODE_NEG) @@ -172,6 +173,7 @@ namespace parser CASE(TOKEN_F16) CASE(TOKEN_F64) CASE(TOKEN_F32) + CASE(TOKEN_BF16) CASE(TOKEN_PRED) CASE(TOKEN_EQ) CASE(TOKEN_NE) @@ -260,6 +262,11 @@ namespace parser CASE(TOKEN_ARRIVE) CASE(TOKEN_RED) CASE(TOKEN_SYNC) + CASE(TOKEN_ALIGNED) + CASE(TOKEN_M16N8K8) + CASE(TOKEN_M16N8K16) + CASE(TOKEN_ROW) + CASE(TOKEN_COL) CASE(TOKEN_POPC) CASE(TOKEN_BALLOT) CASE(TOKEN_F4E) @@ -367,4 +374,3 @@ namespace parser } #endif - diff --git a/ocelot/src/parser/PTXParser.cpp b/ocelot/src/parser/PTXParser.cpp index 636237be..9813d5cb 100644 --- a/ocelot/src/parser/PTXParser.cpp +++ b/ocelot/src/parser/PTXParser.cpp @@ -152,6 +152,56 @@ namespace parser operand.type = instruction.type; } } + //https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#integer-arithmetic-instructions-bfi + //https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#integer-arithmetic-instructions-bfe + if( instruction.opcode == ir::PTXInstruction::Bfi + || instruction.opcode == ir::PTXInstruction::Bfe ) + { + if ( instruction.b.addressMode == ir::PTXOperand::AddressMode::Immediate) + instruction.b.type = ir::PTXOperand::u32; + if ( instruction.c.addressMode == ir::PTXOperand::AddressMode::Immediate) + instruction.c.type = ir::PTXOperand::u32; + } + } + + void PTXParser::State::_setMovVectorImmediateTypes() + { + ir::PTXInstruction& instruction = statement.instruction; + ir::PTXOperand& source = instruction.a; + + if( instruction.opcode != ir::PTXInstruction::Mov || + source.array.empty() ) return; + + unsigned int totalBytes = ir::PTXOperand::bytes( instruction.type ); + if( totalBytes % source.array.size() != 0 ) + { + throw_exception( "Invalid mov vector element size.", + InvalidInstruction ); + } + + unsigned int elementBytes = totalBytes / source.array.size(); + if( source.type == ir::PTXOperand::TypeSpecifier_invalid ) + { + switch( elementBytes ) + { + case 1: source.type = ir::PTXOperand::b8; break; + case 2: source.type = ir::PTXOperand::b16; break; + case 4: source.type = ir::PTXOperand::b32; break; + case 8: source.type = ir::PTXOperand::b64; break; + default: + throw_exception( "Invalid mov vector element size.", + InvalidInstruction ); + } + } + + for( ir::PTXOperand::Array::iterator element = source.array.begin(); + element != source.array.end(); ++element ) + { + if( element->addressMode == ir::PTXOperand::Immediate ) + { + element->type = source.type; + } + } } static std::string strip(const std::string& name) @@ -606,6 +656,8 @@ namespace parser else if( token == TOKEN_SM21 ) statement.targets.push_back( "sm_21" ); else if( token == TOKEN_SM30 ) statement.targets.push_back( "sm_30" ); else if( token == TOKEN_SM35 ) statement.targets.push_back( "sm_35" ); + else if( token == TOKEN_SM50 ) statement.targets.push_back( "sm_50" ); + else if( token == TOKEN_SM86 ) statement.targets.push_back( "sm_86" ); else if( token == TOKEN_MAP_F64_TO_F32 ) { statement.targets.push_back( "map_f64_to_f32" ); @@ -1523,7 +1575,34 @@ namespace parser operandVector.push_back( OperandWrapper( operand, mode->space ) ); } - + void PTXParser::State::vectorOperand( unsigned int elements ) + { + assert( elements == 2 || elements == 4 ); + assert( operandVector.size() >= elements ); + + OperandVector::iterator begin = operandVector.end() - elements; + + ir::PTXOperand vector; + vector.addressMode = ir::PTXOperand::Register; + vector.type = ir::PTXOperand::TypeSpecifier_invalid; + vector.vec = elements == 2 ? ir::PTXOperand::v2 : ir::PTXOperand::v4; + + for( OperandVector::iterator element = begin; + element != operandVector.end(); ++element ) + { + if( vector.type == ir::PTXOperand::TypeSpecifier_invalid && + element->operand.addressMode != ir::PTXOperand::Immediate ) + { + vector.type = element->operand.type; + } + + vector.array.push_back( element->operand ); + } + + operandVector.erase( begin, operandVector.end() ); + operandVector.push_back( vector ); + } + void PTXParser::State::addressableOperand( const std::string& name, long long int value, YYLTYPE& location, bool invert ) { @@ -1843,6 +1922,7 @@ namespace parser } _setImmediateTypes(); + _setMovVectorImmediateTypes(); } void PTXParser::State::instruction( const std::string& opcode ) @@ -1850,6 +1930,35 @@ namespace parser instruction( opcode, TOKEN_B64 ); } + void PTXParser::State::mma( int shapeToken, int accumulatorToken, + int aToken, int bToken, int cToken ) + { + assert( operandVector.size() == 5 ); + ir::PTXOperand::DataType accumulatorType = tokenToDataType( accumulatorToken ); + ir::PTXOperand::DataType aType = tokenToDataType( aToken ); + ir::PTXOperand::DataType bType = tokenToDataType( bToken ); + ir::PTXOperand::DataType cType = tokenToDataType( cToken ); + + statement.directive = ir::PTXStatement::Instr; + statement.instruction.opcode = ir::PTXInstruction::Mma; + statement.instruction.mmaShape = shapeToken == TOKEN_M16N8K8 + ? ir::PTXInstruction::MmaM16N8K8 + : ir::PTXInstruction::MmaM16N8K16; + statement.instruction.type = accumulatorType; + statement.instruction.pg = operandVector[0].operand; + statement.instruction.d = operandVector[1].operand; + statement.instruction.a = operandVector[2].operand; + statement.instruction.b = operandVector[3].operand; + statement.instruction.c = operandVector[4].operand; + + statement.instruction.d.type = accumulatorType; + statement.instruction.c.type = cType; + statement.instruction.a.type = aType; + statement.instruction.b.type = bType; + + _setImmediateTypes(); + } + void PTXParser::State::tex( int dataType ) { report( " Rule: instruction : tex" ); @@ -2455,6 +2564,8 @@ namespace parser case TOKEN_PRED: return ir::PTXOperand::pred; break; case TOKEN_F16: return ir::PTXOperand::f16; break; case TOKEN_F32: return ir::PTXOperand::f32; break; + case TOKEN_BF16: return ir::PTXOperand::bf16; break; + case TOKEN_TF32: return ir::PTXOperand::tf32; break; case TOKEN_F64: return ir::PTXOperand::f64; break; default: { @@ -2533,6 +2644,7 @@ namespace parser if( string == "mad24" ) return ir::PTXInstruction::Mad24; if( string == "mad" ) return ir::PTXInstruction::Mad; if( string == "madc" ) return ir::PTXInstruction::MadC; + if( string == "mma" ) return ir::PTXInstruction::Mma; if( string == "max" ) return ir::PTXInstruction::Max; if( string == "membar" ) return ir::PTXInstruction::Membar; if( string == "min" ) return ir::PTXInstruction::Min; @@ -3000,4 +3112,3 @@ namespace parser } #endif - diff --git a/ocelot/src/parser/ptx.ll b/ocelot/src/parser/ptx.ll index e3d747ef..d58c88d9 100644 --- a/ocelot/src/parser/ptx.ll +++ b/ocelot/src/parser/ptx.ll @@ -166,6 +166,8 @@ LABEL ({IDENTIFIER}{WHITESPACE}":") return OPCODE_MADC; } "mad24" { sstrcpy( yylval->text, yytext, 1024 ); \ return OPCODE_MAD24; } +"mma" { sstrcpy( yylval->text, yytext, 1024 ); \ + return OPCODE_MMA; } "max" { sstrcpy( yylval->text, yytext, 1024 ); \ return OPCODE_MAX; } "mov" { sstrcpy( yylval->text, yytext, 1024 ); \ @@ -333,6 +335,10 @@ LABEL ({IDENTIFIER}{WHITESPACE}":") return TOKEN_SM30; } "sm_35" { yylval->value = TOKEN_SM35; return TOKEN_SM35; } +"sm_50" { yylval->value = TOKEN_SM50; + return TOKEN_SM50; } +"sm_86" { yylval->value = TOKEN_SM86; + return TOKEN_SM86; } "map_f64_to_f32" { yylval->value = TOKEN_MAP_F64_TO_F32; return TOKEN_MAP_F64_TO_F32; } "texmode_independent" { yylval->value = TOKEN_TEXMODE_INDEPENDENT; @@ -355,6 +361,8 @@ LABEL ({IDENTIFIER}{WHITESPACE}":") ".f16" { yylval->value = TOKEN_F16; return TOKEN_F16; } ".f64" { yylval->value = TOKEN_F64; return TOKEN_F64; } ".f32" { yylval->value = TOKEN_F32; return TOKEN_F32; } +".bf16" { yylval->value = TOKEN_BF16; return TOKEN_BF16; } +".tf32" { yylval->value = TOKEN_TF32; return TOKEN_TF32; } ".pred" { yylval->value = TOKEN_PRED; \ return TOKEN_PRED; } @@ -494,6 +502,14 @@ LABEL ({IDENTIFIER}{WHITESPACE}":") return TOKEN_RED; } ".sync" { yylval->value = TOKEN_SYNC; \ return TOKEN_SYNC; } +".aligned" { yylval->value = TOKEN_ALIGNED; \ + return TOKEN_ALIGNED; } +".m16n8k8" { yylval->value = TOKEN_M16N8K8; \ + return TOKEN_M16N8K8; } +".m16n8k16" { yylval->value = TOKEN_M16N8K16; \ + return TOKEN_M16N8K16; } +".row" { yylval->value = TOKEN_ROW; return TOKEN_ROW; } +".col" { yylval->value = TOKEN_COL; return TOKEN_COL; } ".popc" { yylval->value = TOKEN_POPC; \ return TOKEN_POPC; } @@ -648,4 +664,3 @@ void sstrcpy( char* destination, const char* source, unsigned int max ) #endif /******************************************************************************/ - diff --git a/ocelot/src/parser/ptxgrammar.yy b/ocelot/src/parser/ptxgrammar.yy index 32ab8e38..d0d39b56 100644 --- a/ocelot/src/parser/ptxgrammar.yy +++ b/ocelot/src/parser/ptxgrammar.yy @@ -66,6 +66,7 @@ %token OPCODE_POPC OPCODE_PRMT OPCODE_CLZ OPCODE_BFIND OPCODE_BREV %token OPCODE_BFI OPCODE_BFE OPCODE_TESTP OPCODE_TLD4 OPCODE_BAR %token OPCODE_PREFETCH OPCODE_PREFETCHU OPCODE_SHFL OPCODE_SHF +%token OPCODE_MMA %token PREPROCESSOR_INCLUDE PREPROCESSOR_DEFINE PREPROCESSOR_IF %token PREPROCESSOR_IFDEF PREPROCESSOR_ELSE PREPROCESSOR_ENDIF @@ -77,7 +78,7 @@ %token TOKEN_MAXNREG TOKEN_MAXNTID TOKEN_MAXNCTAPERSM TOKEN_MINNCTAPERSM %token TOKEN_SM11 TOKEN_SM12 TOKEN_SM13 TOKEN_SM20 TOKEN_MAP_F64_TO_F32 -%token TOKEN_SM21 TOKEN_SM10 TOKEN_SM30 TOKEN_SM35 +%token TOKEN_SM21 TOKEN_SM10 TOKEN_SM30 TOKEN_SM35 TOKEN_SM50 TOKEN_SM86 %token TOKEN_TEXMODE_INDEPENDENT TOKEN_TEXMODE_UNIFIED %token TOKEN_CONST TOKEN_GLOBAL TOKEN_LOCAL TOKEN_PARAM TOKEN_PRAGMA TOKEN_PTR @@ -86,7 +87,7 @@ %token TOKEN_U32 TOKEN_S32 TOKEN_S8 TOKEN_S16 TOKEN_S64 TOKEN_U8 %token TOKEN_U16 TOKEN_U64 TOKEN_B8 TOKEN_B16 TOKEN_B32 TOKEN_B64 -%token TOKEN_F16 TOKEN_F64 TOKEN_F32 TOKEN_PRED +%token TOKEN_F16 TOKEN_F64 TOKEN_F32 TOKEN_BF16 TOKEN_TF32 TOKEN_PRED %token TOKEN_EQ TOKEN_NE TOKEN_LT TOKEN_LE TOKEN_GT TOKEN_GE %token TOKEN_LS TOKEN_HS TOKEN_EQU TOKEN_NEU TOKEN_LTU TOKEN_LEU @@ -127,7 +128,8 @@ %token TOKEN_TRAP TOKEN_CLAMP TOKEN_ZERO TOKEN_WRAP -%token TOKEN_ARRIVE TOKEN_RED TOKEN_POPC TOKEN_SYNC +%token TOKEN_ARRIVE TOKEN_RED TOKEN_POPC TOKEN_SYNC TOKEN_ALIGNED +%token TOKEN_M16N8K8 TOKEN_M16N8K16 TOKEN_ROW TOKEN_COL %token TOKEN_BALLOT @@ -136,6 +138,8 @@ %token TOKEN_FINITE TOKEN_INFINITE TOKEN_NUMBER TOKEN_NOT_A_NUMBER %token TOKEN_NORMAL TOKEN_SUBNORMAL +%type mmaShape mmaAccumulatorTypeId mmaInputTypeId + %token TOKEN_DECIMAL_CONSTANT %token TOKEN_UNSIGNED_DECIMAL_CONSTANT @@ -260,7 +264,7 @@ singleInitializer : singleList | '{' singleList '}' | '{' singleListSingle '}' | singleListSingle; shaderModel : TOKEN_SM10 | TOKEN_SM11 | TOKEN_SM12 | TOKEN_SM13 | TOKEN_SM20 - | TOKEN_SM21 | TOKEN_SM30 | TOKEN_SM35; + | TOKEN_SM21 | TOKEN_SM30 | TOKEN_SM35 | TOKEN_SM50 | TOKEN_SM86; floatingPointOption : TOKEN_MAP_F64_TO_F32; textureOption: TOKEN_TEXMODE_INDEPENDENT | TOKEN_TEXMODE_UNIFIED; @@ -302,7 +306,8 @@ pointerDataTypeId: TOKEN_U64 | TOKEN_U32; dataTypeId : TOKEN_U8 | TOKEN_U16 | TOKEN_U32 | TOKEN_U64 | TOKEN_S8 | TOKEN_S16 | TOKEN_S32 | TOKEN_S64 | TOKEN_B8 | TOKEN_B16 | TOKEN_B32 - | TOKEN_B64 | TOKEN_F16 | TOKEN_F32 | TOKEN_F64 | TOKEN_PRED; + | TOKEN_B64 | TOKEN_F16 | TOKEN_F32 | TOKEN_F64 + | TOKEN_BF16 | TOKEN_PRED; dataType : dataTypeId { @@ -685,6 +690,19 @@ opcode : OPCODE_COS | OPCODE_SQRT | OPCODE_ADD | OPCODE_RSQRT | OPCODE_ADDC | OPCODE_BFI | OPCODE_TESTP | OPCODE_TLD4 | OPCODE_PREFETCH | OPCODE_PREFETCHU; +mma : OPCODE_MMA TOKEN_SYNC TOKEN_ALIGNED mmaShape TOKEN_ROW TOKEN_COL + mmaAccumulatorTypeId mmaInputTypeId mmaInputTypeId mmaAccumulatorTypeId + arrayOperand ',' arrayOperand ',' arrayOperand ',' arrayOperand ';' +{ + state.mma( $4, $7, $8, $9, $10 ); +}; + +mmaShape : TOKEN_M16N8K8 | TOKEN_M16N8K16; + +mmaAccumulatorTypeId : TOKEN_F16 | TOKEN_F32; + +mmaInputTypeId : TOKEN_F16 | TOKEN_BF16 | TOKEN_TF32; + uninitializableDeclaration : uninitializable addressableVariablePrefix identifier arrayDimensions ';' { @@ -853,7 +871,7 @@ optionalFloatRounding : floatRounding | /* empty string */; instruction : ftzInstruction2 | ftzInstruction3 | approxInstruction2 | basicInstruction3 | bfe | bfi | bfind | brev | branch | addOrSub | addCOrSubC | atom | bar | brkpt | clz | cvt | cvta | isspacep | div | exit - | ld | ldu | mad | mad24 | madc | membar | mov | mul24 | mul | notInstruction + | ld | ldu | mad | mad24 | madc | mma | membar | mov | mul24 | mul | notInstruction | pmevent | popc | prefetch | prefetchu | prmt | rcpSqrtInstruction | red | ret | sad | selp | set | setp | slct | st | suld | suq | sured | sust | testp | tex | tld4 | trap | txq | vote | shfl | shf; @@ -1266,7 +1284,17 @@ movIndexedOperand : identifier '[' TOKEN_DECIMAL_CONSTANT ']' state.indexedOperand( $1, @1, $3 ); }; -movSourceOperand : arrayOperand | offsetAddressableOperand | movIndexedOperand; +movVectorOperand : '{' operand ',' operand '}' +{ + state.vectorOperand(2); +}; + +movVectorOperand : '{' operand ',' operand ',' operand ',' operand '}' +{ + state.vectorOperand(4); +}; + +movSourceOperand : operand | offsetAddressableOperand | movIndexedOperand | movVectorOperand; mov : OPCODE_MOV dataType arrayOperand ',' movSourceOperand ';' {