diff --git a/src/CodeGen_X86.cpp b/src/CodeGen_X86.cpp index 7529c31688f4..e208b0d10981 100644 --- a/src/CodeGen_X86.cpp +++ b/src/CodeGen_X86.cpp @@ -131,6 +131,11 @@ const x86Intrinsic intrinsic_defs[] = { {"llvm.ssub.sat.v16i16", Int(16, 16), "saturating_sub", {Int(16, 16), Int(16, 16)}, Target::AVX2}, {"llvm.ssub.sat.v8i16", Int(16, 8), "saturating_sub", {Int(16, 8), Int(16, 8)}}, + // Sum of absolute differences + {"llvm.x86.sse2.psad.bw", UInt(64, 2), "sum_of_absolute_differences", {UInt(8, 16), UInt(8, 16)}}, + {"llvm.x86.avx2.psad.bw", UInt(64, 4), "sum_of_absolute_differences", {UInt(8, 32), UInt(8, 32)}, Target::AVX2}, + {"llvm.x86.avx512.psad.bw.512", UInt(64, 8), "sum_of_absolute_differences", {UInt(8, 64), UInt(8, 64)}, Target::AVX512_Skylake}, + // Some of the instructions referred to below only appear with // AVX2, but LLVM generates better AVX code if you give it // full 256-bit vectors and let it do the slicing up into @@ -635,10 +640,13 @@ void CodeGen_X86::codegen_vector_reduce(const VectorReduce *op, const Expr &init // One could do a horizontal widening addition with // other dot_products against a vector of ones. Currently disabled // because I haven't found other cases where it's clearly better. - {VectorReduce::Add, 2, u16(wild_u8x_), "horizontal_widening_add", {}, Pattern::SingleArg}, {VectorReduce::Add, 2, i16(wild_u8x_), "horizontal_widening_add", {}, Pattern::SingleArg}, {VectorReduce::Add, 2, i16(wild_i8x_), "horizontal_widening_add", {}, Pattern::SingleArg}, + + // Sum of absolute differences + {VectorReduce::Add, 8, u64(absd(wild_u8x_, wild_u8x_)), "sum_of_absolute_differences", {}}, + }; // clang-format on @@ -708,6 +716,33 @@ void CodeGen_X86::codegen_vector_reduce(const VectorReduce *op, const Expr &init } } + // Rewrite non-native sum-of-absolute-difference variants to the native + // op. We support reducing to various types. We could consider supporting + // multiple reduction factors too, but in general we don't handle non-native + // reduction factors for VectorReduce nodes (yet?). + if (op->op == VectorReduce::Add && + factor == 8) { + const Cast *cast = op->value.as(); + const Call *call = cast ? cast->value.as() : nullptr; + if (call && + call->is_intrinsic(Call::absd) && + cast->type.element_of().can_represent(UInt(8)) && + (cast->type.is_int() || cast->type.is_uint()) && + call->args[0].type().element_of() == UInt(8)) { + + internal_assert(cast->type.element_of() != UInt(64)) << "Should have pattern-matched above\n"; + + // Cast to uint64 instead + Expr equiv = Cast::make(UInt(64, cast->value.type().lanes()), cast->value); + // Reduce on that to hit psadbw + equiv = VectorReduce::make(VectorReduce::Add, equiv, op->type.lanes()); + // Then cast that to the desired type + equiv = Cast::make(cast->type.with_lanes(equiv.type().lanes()), equiv); + codegen(equiv); + return; + } + } + CodeGen_Posix::codegen_vector_reduce(op, init); } diff --git a/test/correctness/simd_op_check.cpp b/test/correctness/simd_op_check.cpp index 5f2258b65f55..ba3e24d6e80c 100644 --- a/test/correctness/simd_op_check.cpp +++ b/test/correctness/simd_op_check.cpp @@ -242,6 +242,18 @@ class SimdOpCheck : public SimdOpCheckTest { check(std::string("packssdw") + check_suffix, 4 * w, i16_sat(i32_1)); check(std::string("packsswb") + check_suffix, 8 * w, i8_sat(i16_1)); check(std::string("packuswb") + check_suffix, 8 * w, u8_sat(i16_1)); + + // Sum-of-absolute-difference ops + { + const int f = 8; // reduction factor. + RDom r(0, f); + check("psadbw", w, sum(u64(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("psadbw", w, sum(u32(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("psadbw", w, sum(u16(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("psadbw", w, sum(i64(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("psadbw", w, sum(i32(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("psadbw", w, sum(i16(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + } } // SSE 3 / SSSE 3 @@ -518,6 +530,18 @@ class SimdOpCheck : public SimdOpCheckTest { check("vpcmpeqq*ymm", 4, select(i64_1 == i64_2, i64(1), i64(2))); check("vpackusdw*ymm", 16, u16(clamp(i32_1, 0, max_u16))); check("vpcmpgtq*ymm", 4, select(i64_1 > i64_2, i64(1), i64(2))); + + // Sum-of-absolute-difference ops + for (int w : {4, 8}) { + const int f = 8; // reduction factor. + RDom r(0, f); + check("vpsadbw", w, sum(u64(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("vpsadbw", w, sum(u32(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("vpsadbw", w, sum(u16(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("vpsadbw", w, sum(i64(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("vpsadbw", w, sum(i32(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + check("vpsadbw", w, sum(i16(absd(in_u8(f * x + r), in_u8(f * x + r + 32))))); + } } if (use_avx512) {