Skip to content

Commit 360b2b2

Browse files
pratiklp00pytorchmergebot
pratiklp00
authored and
pytorchmergebot
committed
re did the operator functions and removed the temporary change from FlashAttentionKernel.cpp
1 parent f336bb0 commit 360b2b2

File tree

11 files changed

+338
-2
lines changed

11 files changed

+338
-2
lines changed

‎aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_double_vsx.h‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -554,6 +554,30 @@ Vectorized<ComplexDbl> inline minimum(
554554
// return _mm256_or_ps(min, isnan);
555555
}
556556

557+
template <>
558+
Vectorized<ComplexDbl> C10_ALWAYS_INLINE operator+(const Vectorized<ComplexDbl>& a, const Vectorized<ComplexDbl>& b) {
559+
return Vectorized<ComplexDbl>{vec_add(a.vec0(), b.vec0()), vec_add(a.vec1(), b.vec1())};
560+
}
561+
562+
template <>
563+
Vectorized<ComplexDbl> C10_ALWAYS_INLINE operator-(const Vectorized<ComplexDbl>& a, const Vectorized<ComplexDbl>& b) {
564+
return Vectorized<ComplexDbl>{vec_sub(a.vec0(), b.vec0()), vec_sub(a.vec1(), b.vec1())};
565+
}
566+
567+
template <>
568+
Vectorized<ComplexDbl> C10_ALWAYS_INLINE operator&(const Vectorized<ComplexDbl>& a, const Vectorized<ComplexDbl>& b) {
569+
return Vectorized<ComplexDbl>{vec_and(a.vec0(), b.vec0()), vec_and(a.vec1(), b.vec1())};
570+
}
571+
572+
template <>
573+
Vectorized<ComplexDbl> C10_ALWAYS_INLINE operator|(const Vectorized<ComplexDbl>& a, const Vectorized<ComplexDbl>& b) {
574+
return Vectorized<ComplexDbl>{vec_or(a.vec0(), b.vec0()), vec_or(a.vec1(), b.vec1())};
575+
}
576+
577+
template <>
578+
Vectorized<ComplexDbl> C10_ALWAYS_INLINE operator^(const Vectorized<ComplexDbl>& a, const Vectorized<ComplexDbl>& b) {
579+
return Vectorized<ComplexDbl>{vec_xor(a.vec0(), b.vec0()), vec_xor(a.vec1(), b.vec1())};
580+
}
557581

558582
} // namespace
559583
} // namespace vec

‎aten/src/ATen/cpu/vec/vec256/vsx/vec256_complex_float_vsx.h‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,13 @@ class Vectorized<ComplexFlt> {
5555
_vec1 = vfloat32{val3.real(), val3.imag(), val4.real(), val4.imag()};
5656
}
5757

58+
C10_ALWAYS_INLINE const vec_internal_type& vec0() const {
59+
return _vec0;
60+
}
61+
C10_ALWAYS_INLINE const vec_internal_type& vec1() const {
62+
return _vec1;
63+
}
64+
5865
template <uint64_t mask>
5966
static std::enable_if_t<blendChoiceComplex(mask) == 0, Vectorized<ComplexFlt>>
6067
C10_ALWAYS_INLINE
@@ -623,6 +630,31 @@ Vectorized<ComplexFlt> inline minimum(
623630
// return _mm256_or_ps(min, isnan);
624631
}
625632

633+
template <>
634+
Vectorized<ComplexFlt> C10_ALWAYS_INLINE operator+(const Vectorized<ComplexFlt>& a, const Vectorized<ComplexFlt>& b) {
635+
return Vectorized<ComplexFlt>{vec_add(a.vec0(), b.vec0()), vec_add(a.vec1(), b.vec1())};
636+
}
637+
638+
template <>
639+
Vectorized<ComplexFlt> C10_ALWAYS_INLINE operator-(const Vectorized<ComplexFlt>& a, const Vectorized<ComplexFlt>& b) {
640+
return Vectorized<ComplexFlt>{vec_sub(a.vec0(), b.vec0()), vec_sub(a.vec1(), b.vec1())};
641+
}
642+
643+
template <>
644+
Vectorized<ComplexFlt> C10_ALWAYS_INLINE operator&(const Vectorized<ComplexFlt>& a, const Vectorized<ComplexFlt>& b) {
645+
return Vectorized<ComplexFlt>{vec_and(a.vec0(), b.vec0()), vec_and(a.vec1(), b.vec1())};
646+
}
647+
648+
template <>
649+
Vectorized<ComplexFlt> C10_ALWAYS_INLINE operator|(const Vectorized<ComplexFlt>& a, const Vectorized<ComplexFlt>& b) {
650+
return Vectorized<ComplexFlt>{vec_or(a.vec0(), b.vec0()), vec_or(a.vec1(), b.vec1())};
651+
}
652+
653+
template <>
654+
Vectorized<ComplexFlt> C10_ALWAYS_INLINE operator^(const Vectorized<ComplexFlt>& a, const Vectorized<ComplexFlt>& b) {
655+
return Vectorized<ComplexFlt>{vec_xor(a.vec0(), b.vec0()), vec_xor(a.vec1(), b.vec1())};
656+
}
657+
626658
} // namespace
627659
} // namespace vec
628660
} // namespace at

‎aten/src/ATen/cpu/vec/vec256/vsx/vec256_double_vsx.h‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -433,6 +433,42 @@ Vectorized<double> inline minimum(
433433
const Vectorized<double>& b) {
434434
return a.minimum(b);
435435
}
436+
437+
template <>
438+
Vectorized<double> C10_ALWAYS_INLINE operator+(const Vectorized<double>& a, const Vectorized<double>& b) {
439+
return Vectorized<double>{vec_add(a.vec0(), b.vec0()), vec_add(a.vec1(), b.vec1())};
440+
}
441+
442+
template <>
443+
Vectorized<double> C10_ALWAYS_INLINE operator-(const Vectorized<double>& a, const Vectorized<double>& b) {
444+
return Vectorized<double>{vec_sub(a.vec0(), b.vec0()), vec_sub(a.vec1(), b.vec1())};
445+
}
446+
447+
template <>
448+
Vectorized<double> C10_ALWAYS_INLINE operator*(const Vectorized<double>& a, const Vectorized<double>& b) {
449+
return Vectorized<double>{vec_mul(a.vec0(), b.vec0()), vec_mul(a.vec1(), b.vec1())};
450+
}
451+
452+
template <>
453+
Vectorized<double> C10_ALWAYS_INLINE operator/(const Vectorized<double>& a, const Vectorized<double>& b) {
454+
return Vectorized<double>{vec_div(a.vec0(), b.vec0()), vec_div(a.vec1(), b.vec1())};
455+
}
456+
457+
template <>
458+
Vectorized<double> C10_ALWAYS_INLINE operator&(const Vectorized<double>& a, const Vectorized<double>& b) {
459+
return Vectorized<double>{vec_and(a.vec0(), b.vec0()), vec_and(a.vec1(), b.vec1())};
460+
}
461+
462+
template <>
463+
Vectorized<double> C10_ALWAYS_INLINE operator|(const Vectorized<double>& a, const Vectorized<double>& b) {
464+
return Vectorized<double>{vec_or(a.vec0(), b.vec0()), vec_or(a.vec1(), b.vec1())};
465+
}
466+
467+
template <>
468+
Vectorized<double> C10_ALWAYS_INLINE operator^(const Vectorized<double>& a, const Vectorized<double>& b) {
469+
return Vectorized<double>{vec_xor(a.vec0(), b.vec0()), vec_xor(a.vec1(), b.vec1())};
470+
}
471+
436472
} // namespace
437473
} // namespace vec
438474
} // namespace at

‎aten/src/ATen/cpu/vec/vec256/vsx/vec256_float_vsx.h‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -456,6 +456,41 @@ Vectorized<float> inline minimum(const Vectorized<float>& a, const Vectorized<fl
456456
return a.minimum(b);
457457
}
458458

459+
template <>
460+
Vectorized<float> C10_ALWAYS_INLINE operator+(const Vectorized<float>& a, const Vectorized<float>& b) {
461+
return Vectorized<float>{vec_add(a.vec0(), b.vec0()), vec_add(a.vec1(), b.vec1())};
462+
}
463+
464+
template <>
465+
Vectorized<float> C10_ALWAYS_INLINE operator-(const Vectorized<float>& a, const Vectorized<float>& b) {
466+
return Vectorized<float>{vec_sub(a.vec0(), b.vec0()), vec_sub(a.vec1(), b.vec1())};
467+
}
468+
469+
template <>
470+
Vectorized<float> C10_ALWAYS_INLINE operator*(const Vectorized<float>& a, const Vectorized<float>& b) {
471+
return Vectorized<float>{vec_mul(a.vec0(), b.vec0()), vec_mul(a.vec1(), b.vec1())};
472+
}
473+
474+
template <>
475+
Vectorized<float> C10_ALWAYS_INLINE operator/(const Vectorized<float>& a, const Vectorized<float>& b) {
476+
return Vectorized<float>{vec_div(a.vec0(), b.vec0()), vec_div(a.vec1(), b.vec1())};
477+
}
478+
479+
template <>
480+
Vectorized<float> C10_ALWAYS_INLINE operator&(const Vectorized<float>& a, const Vectorized<float>& b) {
481+
return Vectorized<float>{vec_and(a.vec0(), b.vec0()), vec_and(a.vec1(), b.vec1())};
482+
}
483+
484+
template <>
485+
Vectorized<float> C10_ALWAYS_INLINE operator|(const Vectorized<float>& a, const Vectorized<float>& b) {
486+
return Vectorized<float>{vec_or(a.vec0(), b.vec0()), vec_or(a.vec1(), b.vec1())};
487+
}
488+
489+
template <>
490+
Vectorized<float> C10_ALWAYS_INLINE operator^(const Vectorized<float>& a, const Vectorized<float>& b) {
491+
return Vectorized<float>{vec_xor(a.vec0(), b.vec0()), vec_xor(a.vec1(), b.vec1())};
492+
}
493+
459494
} // namespace
460495
} // namespace vec
461496
} // namespace at

‎aten/src/ATen/cpu/vec/vec256/vsx/vec256_int16_vsx.h‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,40 @@ Vectorized<int16_t> inline minimum(
362362
return a.minimum(b);
363363
}
364364

365+
template <>
366+
Vectorized<int16_t> C10_ALWAYS_INLINE operator+(const Vectorized<int16_t>& a, const Vectorized<int16_t>& b) {
367+
return Vectorized<int16_t>{vec_add(a.vec0(), b.vec0()), vec_add(a.vec1(), b.vec1())};
368+
}
369+
370+
template <>
371+
Vectorized<int16_t> C10_ALWAYS_INLINE operator-(const Vectorized<int16_t>& a, const Vectorized<int16_t>& b) {
372+
return Vectorized<int16_t>{vec_sub(a.vec0(), b.vec0()), vec_sub(a.vec1(), b.vec1())};
373+
}
374+
375+
template <>
376+
Vectorized<int16_t> C10_ALWAYS_INLINE operator*(const Vectorized<int16_t>& a, const Vectorized<int16_t>& b) {
377+
return Vectorized<int16_t>{vec_mul(a.vec0(), b.vec0()), vec_mul(a.vec1(), b.vec1())};
378+
}
379+
380+
template <>
381+
Vectorized<int16_t> C10_ALWAYS_INLINE operator/(const Vectorized<int16_t>& a, const Vectorized<int16_t>& b) {
382+
return Vectorized<int16_t>{a.vec0()/b.vec0(), a.vec1()/b.vec1()};
383+
}
384+
385+
template <>
386+
Vectorized<int16_t> C10_ALWAYS_INLINE operator&(const Vectorized<int16_t>& a, const Vectorized<int16_t>& b) {
387+
return Vectorized<int16_t>{vec_and(a.vec0(), b.vec0()), vec_and(a.vec1(), b.vec1())};
388+
}
389+
390+
template <>
391+
Vectorized<int16_t> C10_ALWAYS_INLINE operator|(const Vectorized<int16_t>& a, const Vectorized<int16_t>& b) {
392+
return Vectorized<int16_t>{vec_or(a.vec0(), b.vec0()), vec_or(a.vec1(), b.vec1())};
393+
}
394+
395+
template <>
396+
Vectorized<int16_t> C10_ALWAYS_INLINE operator^(const Vectorized<int16_t>& a, const Vectorized<int16_t>& b) {
397+
return Vectorized<int16_t>{vec_xor(a.vec0(), b.vec0()), vec_xor(a.vec1(), b.vec1())};
398+
}
365399

366400
} // namespace
367401
} // namespace vec

‎aten/src/ATen/cpu/vec/vec256/vsx/vec256_int32_vsx.h‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,41 @@ Vectorized<int32_t> inline minimum(
293293
return a.minimum(b);
294294
}
295295

296+
template <>
297+
Vectorized<int32_t> C10_ALWAYS_INLINE operator+(const Vectorized<int32_t>& a, const Vectorized<int32_t>& b) {
298+
return Vectorized<int32_t>{vec_add(a.vec0(), b.vec0()), vec_add(a.vec1(), b.vec1())};
299+
}
300+
301+
template <>
302+
Vectorized<int32_t> C10_ALWAYS_INLINE operator-(const Vectorized<int32_t>& a, const Vectorized<int32_t>& b) {
303+
return Vectorized<int32_t>{vec_sub(a.vec0(), b.vec0()), vec_sub(a.vec1(), b.vec1())};
304+
}
305+
306+
template <>
307+
Vectorized<int32_t> C10_ALWAYS_INLINE operator*(const Vectorized<int32_t>& a, const Vectorized<int32_t>& b) {
308+
return Vectorized<int32_t>{vec_mul(a.vec0(), b.vec0()), vec_mul(a.vec1(), b.vec1())};
309+
}
310+
311+
template <>
312+
Vectorized<int32_t> C10_ALWAYS_INLINE operator/(const Vectorized<int32_t>& a, const Vectorized<int32_t>& b) {
313+
return Vectorized<int32_t>{a.vec0()/b.vec0(), a.vec1()/b.vec1()};
314+
}
315+
316+
template <>
317+
Vectorized<int32_t> C10_ALWAYS_INLINE operator&(const Vectorized<int32_t>& a, const Vectorized<int32_t>& b) {
318+
return Vectorized<int32_t>{vec_and(a.vec0(), b.vec0()), vec_and(a.vec1(), b.vec1())};
319+
}
320+
321+
template <>
322+
Vectorized<int32_t> C10_ALWAYS_INLINE operator|(const Vectorized<int32_t>& a, const Vectorized<int32_t>& b) {
323+
return Vectorized<int32_t>{vec_or(a.vec0(), b.vec0()), vec_or(a.vec1(), b.vec1())};
324+
}
325+
326+
template <>
327+
Vectorized<int32_t> C10_ALWAYS_INLINE operator^(const Vectorized<int32_t>& a, const Vectorized<int32_t>& b) {
328+
return Vectorized<int32_t>{vec_xor(a.vec0(), b.vec0()), vec_xor(a.vec1(), b.vec1())};
329+
}
330+
296331
} // namespace
297332
} // namespace vec
298333
} // namespace at

‎aten/src/ATen/cpu/vec/vec256/vsx/vec256_int64_vsx.h‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,41 @@ Vectorized<int64_t> inline minimum(
246246
return a.minimum(b);
247247
}
248248

249+
template <>
250+
Vectorized<int64_t> C10_ALWAYS_INLINE operator+(const Vectorized<int64_t>& a, const Vectorized<int64_t>& b) {
251+
return Vectorized<int64_t>{vec_add(a.vec0(), b.vec0()), vec_add(a.vec1(), b.vec1())};
252+
}
253+
254+
template <>
255+
Vectorized<int64_t> C10_ALWAYS_INLINE operator-(const Vectorized<int64_t>& a, const Vectorized<int64_t>& b) {
256+
return Vectorized<int64_t>{vec_sub(a.vec0(), b.vec0()), vec_sub(a.vec1(), b.vec1())};
257+
}
258+
259+
template <>
260+
Vectorized<int64_t> C10_ALWAYS_INLINE operator*(const Vectorized<int64_t>& a, const Vectorized<int64_t>& b) {
261+
return Vectorized<int64_t>{vec_mul(a.vec0(), b.vec0()), vec_mul(a.vec1(), b.vec1())};
262+
}
263+
264+
template <>
265+
Vectorized<int64_t> C10_ALWAYS_INLINE operator/(const Vectorized<int64_t>& a, const Vectorized<int64_t>& b) {
266+
return Vectorized<int64_t>{vec_div(a.vec0(), b.vec0()), vec_div(a.vec1(), b.vec1())};
267+
}
268+
269+
template <>
270+
Vectorized<int64_t> C10_ALWAYS_INLINE operator&(const Vectorized<int64_t>& a, const Vectorized<int64_t>& b) {
271+
return Vectorized<int64_t>{vec_and(a.vec0(), b.vec0()), vec_and(a.vec1(), b.vec1())};
272+
}
273+
274+
template <>
275+
Vectorized<int64_t> C10_ALWAYS_INLINE operator|(const Vectorized<int64_t>& a, const Vectorized<int64_t>& b) {
276+
return Vectorized<int64_t>{vec_or(a.vec0(), b.vec0()), vec_or(a.vec1(), b.vec1())};
277+
}
278+
279+
template <>
280+
Vectorized<int64_t> C10_ALWAYS_INLINE operator^(const Vectorized<int64_t>& a, const Vectorized<int64_t>& b) {
281+
return Vectorized<int64_t>{vec_xor(a.vec0(), b.vec0()), vec_xor(a.vec1(), b.vec1())};
282+
}
283+
249284
} // namespace
250285
} // namespace vec
251286
} // namespace at

‎aten/src/ATen/cpu/vec/vec256/vsx/vec256_qint32_vsx.h‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,42 @@ Vectorized<c10::qint32> inline minimum(
240240
const Vectorized<c10::qint32>& b) {
241241
return a.minimum(b);
242242
}
243+
244+
template <>
245+
Vectorized<c10::qint32> C10_ALWAYS_INLINE operator+(const Vectorized<c10::qint32>& a, const Vectorized<c10::qint32>& b) {
246+
return Vectorized<c10::qint32>{vec_add(a.vec0(), b.vec0()), vec_add(a.vec1(), b.vec1())};
247+
}
248+
249+
template <>
250+
Vectorized<c10::qint32> C10_ALWAYS_INLINE operator-(const Vectorized<c10::qint32>& a, const Vectorized<c10::qint32>& b) {
251+
return Vectorized<c10::qint32>{vec_sub(a.vec0(), b.vec0()), vec_sub(a.vec1(), b.vec1())};
252+
}
253+
254+
template <>
255+
Vectorized<c10::qint32> C10_ALWAYS_INLINE operator*(const Vectorized<c10::qint32>& a, const Vectorized<c10::qint32>& b) {
256+
return Vectorized<c10::qint32>{vec_mul(a.vec0(), b.vec0()), vec_mul(a.vec1(), b.vec1())};
257+
}
258+
259+
template <>
260+
Vectorized<c10::qint32> C10_ALWAYS_INLINE operator/(const Vectorized<c10::qint32>& a, const Vectorized<c10::qint32>& b) {
261+
return Vectorized<c10::qint32>{a.vec0()/b.vec0(), a.vec1()/b.vec1()};
262+
}
263+
264+
template <>
265+
Vectorized<c10::qint32> C10_ALWAYS_INLINE operator&(const Vectorized<c10::qint32>& a, const Vectorized<c10::qint32>& b) {
266+
return Vectorized<c10::qint32>{vec_and(a.vec0(), b.vec0()), vec_and(a.vec1(), b.vec1())};
267+
}
268+
269+
template <>
270+
Vectorized<c10::qint32> C10_ALWAYS_INLINE operator|(const Vectorized<c10::qint32>& a, const Vectorized<c10::qint32>& b) {
271+
return Vectorized<c10::qint32>{vec_or(a.vec0(), b.vec0()), vec_or(a.vec1(), b.vec1())};
272+
}
273+
274+
template <>
275+
Vectorized<c10::qint32> C10_ALWAYS_INLINE operator^(const Vectorized<c10::qint32>& a, const Vectorized<c10::qint32>& b) {
276+
return Vectorized<c10::qint32>{vec_xor(a.vec0(), b.vec0()), vec_xor(a.vec1(), b.vec1())};
277+
}
278+
243279
} // namespace
244280
} // namespace vec
245281
} // namespace at

‎aten/src/ATen/cpu/vec/vec256/vsx/vec256_qint8_vsx.h‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -442,6 +442,42 @@ Vectorized<c10::qint8> inline minimum(
442442
const Vectorized<c10::qint8>& b) {
443443
return a.minimum(b);
444444
}
445+
446+
template <>
447+
Vectorized<c10::qint8> C10_ALWAYS_INLINE operator+(const Vectorized<c10::qint8>& a, const Vectorized<c10::qint8>& b) {
448+
return Vectorized<c10::qint8>{vec_add(a.vec0(), b.vec0()), vec_add(a.vec1(), b.vec1())};
449+
}
450+
451+
template <>
452+
Vectorized<c10::qint8> C10_ALWAYS_INLINE operator-(const Vectorized<c10::qint8>& a, const Vectorized<c10::qint8>& b) {
453+
return Vectorized<c10::qint8>{vec_sub(a.vec0(), b.vec0()), vec_sub(a.vec1(), b.vec1())};
454+
}
455+
456+
template <>
457+
Vectorized<c10::qint8> C10_ALWAYS_INLINE operator*(const Vectorized<c10::qint8>& a, const Vectorized<c10::qint8>& b) {
458+
return Vectorized<c10::qint8>{vec_mul(a.vec0(), b.vec0()), vec_mul(a.vec1(), b.vec1())};
459+
}
460+
461+
template <>
462+
Vectorized<c10::qint8> C10_ALWAYS_INLINE operator/(const Vectorized<c10::qint8>& a, const Vectorized<c10::qint8>& b) {
463+
return Vectorized<c10::qint8>{a.vec0()/b.vec0(), a.vec1()/b.vec1()};
464+
}
465+
466+
template <>
467+
Vectorized<c10::qint8> C10_ALWAYS_INLINE operator&(const Vectorized<c10::qint8>& a, const Vectorized<c10::qint8>& b) {
468+
return Vectorized<c10::qint8>{vec_and(a.vec0(), b.vec0()), vec_and(a.vec1(), b.vec1())};
469+
}
470+
471+
template <>
472+
Vectorized<c10::qint8> C10_ALWAYS_INLINE operator|(const Vectorized<c10::qint8>& a, const Vectorized<c10::qint8>& b) {
473+
return Vectorized<c10::qint8>{vec_or(a.vec0(), b.vec0()), vec_or(a.vec1(), b.vec1())};
474+
}
475+
476+
template <>
477+
Vectorized<c10::qint8> C10_ALWAYS_INLINE operator^(const Vectorized<c10::qint8>& a, const Vectorized<c10::qint8>& b) {
478+
return Vectorized<c10::qint8>{vec_xor(a.vec0(), b.vec0()), vec_xor(a.vec1(), b.vec1())};
479+
}
480+
445481
} // namespace
446482
} // namespace vec
447483
} // namespace at

0 commit comments

Comments
 (0)