File tree 2 files changed +20
-0
lines changed
2 files changed +20
-0
lines changed Original file line number Diff line number Diff line change @@ -133,6 +133,16 @@ class BasicBlockSimplify : public IRVisitor {
133
133
}
134
134
}
135
135
136
+ void visit (UnaryOpStmt *stmt) override {
137
+ if (stmt->op_type == UnaryOpType::abs ) {
138
+ auto operand_type = stmt->operand ->ret_type ;
139
+ if (is_integral (operand_type) && is_unsigned (operand_type)) {
140
+ // abs(u) -> u
141
+ stmt->replace_usages_with (stmt->operand );
142
+ modifier.erase (stmt);
143
+ }
144
+ }
145
+ }
136
146
template <typename T>
137
147
static bool identical_vectors (const std::vector<T> &a,
138
148
const std::vector<T> &b) {
Original file line number Diff line number Diff line change @@ -78,3 +78,13 @@ def foo(x: ti.i64) -> ti.i64:
78
78
79
79
for x in [- (2 ** 40 ), 0 , 2 ** 40 ]:
80
80
assert foo (x ) == abs (x )
81
+
82
+
83
+ @test_utils .test ()
84
+ def test_abs_u32 ():
85
+ @ti .kernel
86
+ def foo (x : ti .u32 ) -> ti .u32 :
87
+ return abs (x )
88
+
89
+ for x in [0 , 2 ** 20 ]:
90
+ assert foo (x ) == abs (x )
You can’t perform that action at this time.
0 commit comments