Skip to content

Commit ea2523b

Browse files
avik-palavik-pal
committed
fix: be conservative if effect getValue is null
1 parent 38b1842 commit ea2523b

File tree

3 files changed

+89
-8
lines changed

3 files changed

+89
-8
lines changed

‎src/enzyme_ad/jax/Passes/MarkFunctionMemoryEffectsPass.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,8 @@ struct MarkFunctionMemoryEffectsPass
189189
auto &memEffects = memEffectsOrNothing.value();
190190

191191
for (const auto &effect : memEffects) {
192-
if (effect.getValue() && effect.getValue() == operand->get()) {
192+
if (!effect.getValue() ||
193+
(effect.getValue() && effect.getValue() == operand->get())) {
193194
if (isa<MemoryEffects::Read>(effect.getEffect())) {
194195
effects.set(0);
195196
} else if (isa<MemoryEffects::Write>(effect.getEffect())) {

‎test/lit_tests/memoryeffects/add_kernel.mlir‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
// RUN: enzymexlamlir-opt --mark-func-memory-effects %s | FileCheck %s
2+
3+
module {
4+
// CHECK: llvm.func @add_kernel(%arg0: !llvm.ptr<1> {enzymexla.memory_effects = ["read", "write"], llvm.nofree}, %arg1: !llvm.ptr<1> {enzymexla.memory_effects = ["read", "write"], llvm.nofree}, %arg2: !llvm.ptr<1> {enzymexla.memory_effects = ["read", "write"], llvm.nofree}, %arg3: !llvm.ptr<1> {enzymexla.memory_effects = [], llvm.nofree}) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"], noinline = false, nvvm.kernel = 1 : ui1, nvvm.reqntid = array<i32: 32>, ttg.global_scratch_memory_alignment = 1 : i32, ttg.global_scratch_memory_size = 0 : i32} {
5+
llvm.func @add_kernel(%arg0: !llvm.ptr<1>, %arg1: !llvm.ptr<1>, %arg2: !llvm.ptr<1>, %arg3: !llvm.ptr<1>) attributes {noinline = false, nvvm.kernel = 1 : ui1, nvvm.reqntid = array<i32: 32>, ttg.global_scratch_memory_alignment = 1 : i32, ttg.global_scratch_memory_size = 0 : i32} {
6+
%0 = llvm.mlir.undef : vector<1xf32>
7+
%1 = llvm.mlir.constant(0 : i32) : i32
8+
%2 = llvm.mlir.constant(32 : i32) : i32
9+
%3 = llvm.mlir.constant(31 : i32) : i32
10+
%4 = llvm.mlir.constant(0 : index) : i32
11+
%5 = llvm.mlir.constant(1024 : i32) : i32
12+
%6 = llvm.mlir.constant(64 : i32) : i32
13+
%7 = llvm.call_intrinsic "llvm.nvvm.read.ptx.sreg.ctaid.x"() : () -> i32
14+
%8 = llvm.mul %7, %6 : i32
15+
%9 = nvvm.read.ptx.sreg.tid.x : i32
16+
%10 = llvm.and %9, %3 : i32
17+
%11 = llvm.shl %10, %1 : i32
18+
%12 = llvm.or %1, %11 : i32
19+
%13 = llvm.or %12, %1 : i32
20+
%14 = llvm.and %13, %3 : i32
21+
%15 = llvm.lshr %14, %1 : i32
22+
%16 = llvm.xor %1, %15 : i32
23+
%17 = llvm.xor %1, %16 : i32
24+
%18 = llvm.xor %17, %1 : i32
25+
%19 = llvm.xor %17, %2 : i32
26+
%20 = llvm.add %18, %4 : i32
27+
%21 = llvm.add %19, %4 : i32
28+
%22 = llvm.add %8, %20 : i32
29+
%23 = llvm.add %8, %21 : i32
30+
%24 = llvm.icmp "slt" %22, %5 : i32
31+
%25 = llvm.icmp "slt" %23, %5 : i32
32+
%26 = llvm.getelementptr %arg0[%22] : (!llvm.ptr<1>, i32) -> !llvm.ptr<1>, f32
33+
%27 = llvm.getelementptr %arg0[%23] : (!llvm.ptr<1>, i32) -> !llvm.ptr<1>, f32
34+
%28 = llvm.inline_asm has_side_effects asm_dialect = att operand_attrs = [] "mov.u32 $0, 0x0;\0A\09@$2 ld.global.b32 { $0 }, [ $1 + 0 ];", "=r,l,b" %26, %24 : (!llvm.ptr<1>, i1) -> i32
35+
%29 = llvm.bitcast %28 : i32 to vector<1xf32>
36+
%30 = llvm.extractelement %29[%4 : i32] : vector<1xf32>
37+
%31 = llvm.inline_asm has_side_effects asm_dialect = att operand_attrs = [] "mov.u32 $0, 0x0;\0A\09@$2 ld.global.b32 { $0 }, [ $1 + 0 ];", "=r,l,b" %27, %25 : (!llvm.ptr<1>, i1) -> i32
38+
%32 = llvm.bitcast %31 : i32 to vector<1xf32>
39+
%33 = llvm.extractelement %32[%4 : i32] : vector<1xf32>
40+
%34 = llvm.getelementptr %arg1[%22] : (!llvm.ptr<1>, i32) -> !llvm.ptr<1>, f32
41+
%35 = llvm.getelementptr %arg1[%23] : (!llvm.ptr<1>, i32) -> !llvm.ptr<1>, f32
42+
%36 = llvm.inline_asm has_side_effects asm_dialect = att operand_attrs = [] "mov.u32 $0, 0x0;\0A\09@$2 ld.global.b32 { $0 }, [ $1 + 0 ];", "=r,l,b" %34, %24 : (!llvm.ptr<1>, i1) -> i32
43+
%37 = llvm.bitcast %36 : i32 to vector<1xf32>
44+
%38 = llvm.extractelement %37[%4 : i32] : vector<1xf32>
45+
%39 = llvm.inline_asm has_side_effects asm_dialect = att operand_attrs = [] "mov.u32 $0, 0x0;\0A\09@$2 ld.global.b32 { $0 }, [ $1 + 0 ];", "=r,l,b" %35, %25 : (!llvm.ptr<1>, i1) -> i32
46+
%40 = llvm.bitcast %39 : i32 to vector<1xf32>
47+
%41 = llvm.extractelement %40[%4 : i32] : vector<1xf32>
48+
%42 = llvm.fadd %30, %38 : f32
49+
%43 = llvm.fadd %33, %41 : f32
50+
%44 = llvm.getelementptr %arg2[%22] : (!llvm.ptr<1>, i32) -> !llvm.ptr<1>, f32
51+
%45 = llvm.getelementptr %arg2[%23] : (!llvm.ptr<1>, i32) -> !llvm.ptr<1>, f32
52+
%46 = llvm.insertelement %42, %0[%1 : i32] : vector<1xf32>
53+
%47 = llvm.bitcast %46 : vector<1xf32> to i32
54+
%48 = llvm.inline_asm has_side_effects asm_dialect = att operand_attrs = [] "@$2 st.global.b32 [ $1 + 0 ], { $0 };", "r,l,b" %47, %44, %24 : (i32, !llvm.ptr<1>, i1) -> !llvm.void
55+
%49 = llvm.insertelement %43, %0[%1 : i32] : vector<1xf32>
56+
%50 = llvm.bitcast %49 : vector<1xf32> to i32
57+
%51 = llvm.inline_asm has_side_effects asm_dialect = att operand_attrs = [] "@$2 st.global.b32 [ $1 + 0 ], { $0 };", "r,l,b" %50, %45, %25 : (i32, !llvm.ptr<1>, i1) -> !llvm.void
58+
llvm.return
59+
}
60+
}
61+
62+
module {
63+
// CHECK: tt.func @add_kernel_call(%arg0: !tt.ptr<f32> {enzymexla.memory_effects = ["read"], llvm.nofree, llvm.readonly}, %arg1: !tt.ptr<f32> {enzymexla.memory_effects = ["read"], llvm.nofree, llvm.readonly}, %arg2: !tt.ptr<f32> {enzymexla.memory_effects = ["write"], llvm.nofree, llvm.writeonly}) attributes {enzymexla.memory_effects = ["read", "write"], noinline = false}
64+
tt.func @add_kernel_call(%arg0: !tt.ptr<f32>, %arg1: !tt.ptr<f32>, %arg2: !tt.ptr<f32>) attributes {noinline = false} {
65+
%cst = arith.constant dense<1024> : tensor<64xi32>
66+
%c64_i32 = arith.constant 64 : i32
67+
%0 = tt.get_program_id x : i32
68+
%1 = arith.muli %0, %c64_i32 : i32
69+
%2 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
70+
%3 = tt.splat %1 : i32 -> tensor<64xi32>
71+
%4 = arith.addi %3, %2 : tensor<64xi32>
72+
%5 = arith.cmpi slt, %4, %cst : tensor<64xi32>
73+
%6 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<64x!tt.ptr<f32>>
74+
%7 = tt.addptr %6, %4 : tensor<64x!tt.ptr<f32>>, tensor<64xi32>
75+
%8 = tt.load %7, %5 : tensor<64x!tt.ptr<f32>>
76+
%9 = tt.splat %arg1 : !tt.ptr<f32> -> tensor<64x!tt.ptr<f32>>
77+
%10 = tt.addptr %9, %4 : tensor<64x!tt.ptr<f32>>, tensor<64xi32>
78+
%11 = tt.load %10, %5 : tensor<64x!tt.ptr<f32>>
79+
%12 = arith.addf %8, %11 : tensor<64xf32>
80+
%13 = tt.splat %arg2 : !tt.ptr<f32> -> tensor<64x!tt.ptr<f32>>
81+
%14 = tt.addptr %13, %4 : tensor<64x!tt.ptr<f32>>, tensor<64xi32>
82+
tt.store %14, %12, %5 : tensor<64x!tt.ptr<f32>>
83+
tt.return
84+
}
85+
}

‎test/lit_tests/memoryeffects/shlo.mlir‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,8 @@ module {
3434
return %0 : tensor<64xi64>
3535
}
3636

37-
// ASSUME: @main4(%arg0: tensor<64xi64> {enzymexla.memory_effects = []}) -> tensor<64xi64> attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
38-
// NOASSUME: func.func @main4(%arg0: tensor<64xi64> {enzymexla.memory_effects = []}) -> tensor<64xi64> attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
37+
// ASSUME: @main4(%arg0: tensor<64xi64> {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}) -> tensor<64xi64> attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
38+
// NOASSUME: func.func @main4(%arg0: tensor<64xi64> {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}) -> tensor<64xi64> attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
3939

4040
func.func @main5(%arg0: tensor<64xi64>) -> tensor<64xi64> {
4141
%0 = stablehlo.custom_call @mycall1(%arg0) {has_side_effect = false} : (tensor<64xi64>) -> tensor<64xi64>
@@ -44,9 +44,4 @@ module {
4444

4545
// ASSUME: @main5(%arg0: tensor<64xi64> {enzymexla.memory_effects = []}) -> tensor<64xi64> attributes {enzymexla.memory_effects = []} {
4646
// NOASSUME: func.func @main5(%arg0: tensor<64xi64> {enzymexla.memory_effects = []}) -> tensor<64xi64> attributes {enzymexla.memory_effects = []} {
47-
48-
4947
}
50-
51-
52-

0 commit comments

Comments
 (0)