From a029585e0a036199b672b42ac46cfd252fc706cb Mon Sep 17 00:00:00 2001 From: Christian Sigg Date: Wed, 28 Aug 2024 19:21:41 +0200 Subject: [PATCH] Fix underflow in highestPowOf2Divisor() (#4562) Prevent highestPowOf2Divisor() from producing an underflow when the argument is INT_MIN. Co-authored-by: moerafaat --- include/triton/Dialect/Triton/IR/Utility.h | 6 +++++- test/Analysis/test-alignment.mlir | 11 +++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/include/triton/Dialect/Triton/IR/Utility.h b/include/triton/Dialect/Triton/IR/Utility.h index 0ef5971473..1ff63697ec 100644 --- a/include/triton/Dialect/Triton/IR/Utility.h +++ b/include/triton/Dialect/Triton/IR/Utility.h @@ -31,7 +31,11 @@ template Int ceil(Int m, Int n) { return (m + n - 1) / n; } /// Get the highest power of 2 divisor of an integer. template T highestPowOf2Divisor(T n) { - if (n == 0) { + // When n is 0 or min, return the highest power of 2. The min case is handled + // separately to avoid underflow when T is a signed integer. Technically + // in that case the correct divisor is -n, but this value is outside the + // range of possible values, so we take the next best alternative. + if (n == 0 || n == std::numeric_limits::min()) { return (static_cast(1) << (sizeof(T) * 8 - 2)); } return (n & (~(n - 1))); diff --git a/test/Analysis/test-alignment.mlir b/test/Analysis/test-alignment.mlir index 4161c08ad5..d2664b959d 100644 --- a/test/Analysis/test-alignment.mlir +++ b/test/Analysis/test-alignment.mlir @@ -863,3 +863,14 @@ tt.func public @chained_for(%8: tensor<128x64x!tt.ptr> {tt.divisibility = } tt.return } + +// ----- + +// CHECK-LABEL: @int_min_does_not_underflow_in_analysis +module { + tt.func @int_min_does_not_underflow_in_analysis() -> i64 { + // CHECK: divisibility = [4611686018427387904] + %int_min = arith.constant -9223372036854775808 : i64 + tt.return %int_min : i64 + } +}