From 800b54b09bbde55c98407de478ab1508f0ba2f8f Mon Sep 17 00:00:00 2001 From: Ahan Gupta Date: Fri, 30 Jan 2026 14:03:16 -0600 Subject: [PATCH 1/2] add buggy softmax fwd mode test --- test/lit_tests/diffrules/triton/softmax.mlir | 90 ++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 test/lit_tests/diffrules/triton/softmax.mlir diff --git a/test/lit_tests/diffrules/triton/softmax.mlir b/test/lit_tests/diffrules/triton/softmax.mlir new file mode 100644 index 0000000000..692ffa9416 --- /dev/null +++ b/test/lit_tests/diffrules/triton/softmax.mlir @@ -0,0 +1,90 @@ +// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=softmax_kernel outfn= argTys=enzyme_dup,enzyme_dup,enzyme_const,enzyme_const,enzyme_const,enzyme_const retTys=enzyme_dup mode=ForwardMode" --canonicalize | FileCheck %s + +module { + enzymexla_tt_ext.module @softmax_kernel_tt { + builtin.module @softmax_kernel_inner { + tt.func public @softmax_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32, %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<0xFF800000> : tensor<64xf32> + %0 = tt.get_program_id x : i32 + %1 = tt.get_num_programs x : i32 + scf.for %arg6 = %0 to %arg4 step %1 : i32 { + %2 = arith.muli %arg6, %arg2 : i32 + %3 = tt.addptr %arg1, %2 : !tt.ptr, i32 + %4 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %5 = tt.splat %3 : !tt.ptr -> tensor<64x!tt.ptr> + %6 = tt.addptr %5, %4 : tensor<64x!tt.ptr>, tensor<64xi32> + %7 = tt.splat %arg5 : i32 -> tensor<64xi32> + %8 = arith.cmpi slt, %4, %7 : tensor<64xi32> + %9 = tt.load %6, %8, %cst : tensor<64x!tt.ptr> + %10 = "tt.reduce"(%9) <{axis = 0 : i32}> ({ + ^bb0(%arg7: f32 , %arg8: f32): + %21 = arith.maxnumf %arg7, %arg8 : f32 + tt.reduce.return %21 : f32 + }) : (tensor<64xf32>) -> f32 + %11 = tt.splat %10 : f32 -> tensor<64xf32> + %12 = arith.subf %9, %11 : tensor<64xf32> + %13 = math.exp %12 : tensor<64xf32> + %14 = "tt.reduce"(%13) <{axis = 0 : i32}> ({ + ^bb0(%arg7: f32, %arg8: f32): + %21 = arith.addf %arg7, %arg8 : f32 + tt.reduce.return %21 : f32 + }) : (tensor<64xf32>) -> f32 + %15 = tt.splat %14 : f32 -> tensor<64xf32> + %16 = arith.divf %13, %15 : tensor<64xf32> + %17 = arith.muli %arg6, %arg3 : i32 + %18 = tt.addptr %arg0, %17 : !tt.ptr, i32 + %19 = tt.splat %18 : !tt.ptr -> tensor<64x!tt.ptr> + %20 = tt.addptr %19, %4 : tensor<64x!tt.ptr>, tensor<64xi32> + tt.store %20, %16, %8 : tensor<64x!tt.ptr> + } {tt.num_stages = 2 : i32} + tt.return + } + } + } + func.func @main(%arg0: tensor<1024xf32>, %arg1: tensor<1024xf32>, %arg2: tensor<1024xf32>, %arg3: tensor) -> (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>) { + %c_0 = stablehlo.constant dense<1> : tensor + %c_1 = stablehlo.constant dense<16> : tensor + %0:3 = enzymexla_tt_ext.call @softmax_kernel_tt::@softmax_kernel_inner::@softmax_kernel clusters in (%c_0, %c_0, %c_0) blocks in(%c_1, %c_0, %c_0) (%arg0, %arg1, %arg3) {output_operand_aliases = [#stablehlo.output_operand_alias, #stablehlo.output_operand_alias]} : (tensor<1024xf32>, tensor<1024xf32>, tensor) -> (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>) + return %0#0, %0#1, %0#2 : tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32> + } +} + + +// CHECK: tt.func public @fwddiffsoftmax_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: !tt.ptr {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: !tt.ptr {tt.divisibility = 16 : i32}, %arg6: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { +// CHECK-NEXT: %cst = arith.constant dense<0xFF800000> : tensor<64xf32> loc(#loc1) +// CHECK-NEXT: %0 = tt.get_program_id x : i32 loc(#loc2) +// CHECK-NEXT: %[[v0:.+]] = tt.get_program_id x : i32 loc(#loc2) +// CHECK-NEXT: %1 = tt.get_num_programs x : i32 loc(#loc3) +// CHECK-NEXT: %[[v1:.+]] = tt.get_num_programs x : i32 loc(#loc3) +// CHECK-NEXT: scf.for %arg6 = %[[v0:.+]] to %arg4 step %[[v1:.+]] : i32 { +// CHECK-NEXT: %2 = arith.muli %arg6, %arg2 : i32 loc(#loc5) +// CHECK-NEXT: %3 = tt.addptr %arg1, %2 : !tt.ptr, i32 loc(#loc6) +// CHECK-NEXT: %4 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> loc(#loc7) +// CHECK-NEXT: %5 = tt.splat %3 : !tt.ptr -> tensor<64x!tt.ptr> loc(#loc8) +// CHECK-NEXT: %6 = tt.addptr %5, %4 : tensor<64x!tt.ptr>, tensor<64xi32> loc(#loc8) +// CHECK-NEXT: %7 = tt.splat %arg5 : i32 -> tensor<64xi32> loc(#loc9) +// CHECK-NEXT: %8 = arith.cmpi slt, %4, %7 : tensor<64xi32> loc(#loc9) +// CHECK-NEXT: %9 = tt.load %6, %8, %cst : tensor<64x!tt.ptr> loc(#loc10) +// CHECK-NEXT: %10 = "tt.reduce"(%9) <{axis = 0 : i32}> ({ +// CHECK-NEXT: ^bb0(%arg7: f32 loc(callsite(#loc1 at #loc12)), %arg8: f32 loc(callsite(#loc1 at #loc12))): +// CHECK-NEXT: %21 = arith.maxnumf %arg7, %arg8 : f32 loc(#loc29) +// CHECK-NEXT: tt.reduce.return %21 : f32 loc(#loc25) +// CHECK-NEXT: }) : (tensor<64xf32>) -> f32 loc(#loc25) +// CHECK-NEXT: %11 = tt.splat %10 : f32 -> tensor<64xf32> loc(#loc14) +// CHECK-NEXT: %12 = arith.subf %9, %11 : tensor<64xf32> loc(#loc14) +// CHECK-NEXT: %13 = math.exp %12 : tensor<64xf32> loc(#loc15) +// CHECK-NEXT: %14 = "tt.reduce"(%13) <{axis = 0 : i32}> ({ +// CHECK-NEXT: ^bb0(%arg7: f32 loc(callsite(#loc1 at #loc17)), %arg8: f32 loc(callsite(#loc1 at #loc17))): +// CHECK-NEXT: %21 = arith.addf %arg7, %arg8 : f32 loc(#loc30) +// CHECK-NEXT: tt.reduce.return %21 : f32 loc(#loc27) +// CHECK-NEXT: }) : (tensor<64xf32>) -> f32 loc(#loc27) +// CHECK-NEXT: %15 = tt.splat %14 : f32 -> tensor<64xf32> loc(#loc19) +// CHECK-NEXT: %16 = arith.divf %13, %15 : tensor<64xf32> loc(#loc19) +// CHECK-NEXT: %17 = arith.muli %arg6, %arg3 : i32 loc(#loc20) +// CHECK-NEXT: %18 = tt.addptr %arg0, %17 : !tt.ptr, i32 loc(#loc21) +// CHECK-NEXT: %19 = tt.splat %18 : !tt.ptr -> tensor<64x!tt.ptr> loc(#loc22) +// CHECK-NEXT: %20 = tt.addptr %19, %4 : tensor<64x!tt.ptr>, tensor<64xi32> loc(#loc22) +// CHECK-NEXT: tt.store %20, %16, %8 : tensor<64x!tt.ptr> loc(#loc23) +// CHECK-NEXT: } {tt.num_stages = 2 : i32} loc(#loc4) +// CHECK-NEXT: tt.return loc(#loc24) +// CHECK-NEXT: } From c4231b678066910fd03d53fd1ff4465288a08c0a Mon Sep 17 00:00:00 2001 From: Ahan Gupta Date: Fri, 5 Jun 2026 18:30:12 -0500 Subject: [PATCH 2/2] add fwd mode rules + relevant testing --- .../TritonAutoDiffOpInterfaceImpl.cpp | 202 ++++++++++++++++++ .../jax/Implementations/TritonDerivatives.td | 10 + test/lit_tests/diffrules/triton/matmul.mlir | 115 ++++++++++ .../diffrules/triton/op_coverage.mlir | 85 ++++++++ test/lit_tests/diffrules/triton/softmax.mlir | 66 ++---- test/lit_tests/triton/matmul_kernel.mlir | 37 ++++ test/lit_tests/triton/softmax_kernel.mlir | 37 ++++ 7 files changed, 506 insertions(+), 46 deletions(-) create mode 100644 test/lit_tests/diffrules/triton/matmul.mlir create mode 100644 test/lit_tests/diffrules/triton/op_coverage.mlir create mode 100644 test/lit_tests/triton/matmul_kernel.mlir create mode 100644 test/lit_tests/triton/softmax_kernel.mlir diff --git a/src/enzyme_ad/jax/Implementations/TritonAutoDiffOpInterfaceImpl.cpp b/src/enzyme_ad/jax/Implementations/TritonAutoDiffOpInterfaceImpl.cpp index e9071f1798..6385756d2d 100644 --- a/src/enzyme_ad/jax/Implementations/TritonAutoDiffOpInterfaceImpl.cpp +++ b/src/enzyme_ad/jax/Implementations/TritonAutoDiffOpInterfaceImpl.cpp @@ -13,9 +13,13 @@ #include "Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.h" #include "Enzyme/MLIR/Interfaces/AutoDiffOpInterface.h" +#include "Enzyme/MLIR/Interfaces/AutoDiffTypeInterface.h" #include "Enzyme/MLIR/Interfaces/GradientUtils.h" #include "Enzyme/MLIR/Interfaces/GradientUtilsReverse.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/IRMapping.h" + #include "triton/Dialect/Triton/IR/Dialect.h" #include "src/enzyme_ad/jax/Implementations/XLADerivatives.h" @@ -86,6 +90,201 @@ class AutoDiffTritonFuncFunctionInterface } }; +/// Forward mode autodiff for Triton ReduceOp. +/// +/// Triton reduce syntax: +/// %out = "tt.reduce"(%input) <{axis = 0 : i32}> ({ +/// ^bb0(%arg0: f32, %arg1: f32): +/// %res = arith.addf %arg0, %arg1 : f32 +/// tt.reduce.return %res : f32 +/// }) : (tensor<64xf32>) -> f32 +/// +/// Key differences from StableHLO: +/// - No init_value operand (Triton infers identity from combiner) +/// - Block args are scalars (f32), not rank-0 tensors +/// - Single axis attribute, not array of dimensions +/// - Result can be scalar when reducing all dimensions +/// +/// For sum reduction: d(sum(x)) = sum(dx) +/// For max/min: derivative flows through the selected element +class AutoDiffTritonReduceFwd + : public AutoDiffOpInterface::ExternalModel { +public: + LogicalResult createForwardModeTangent(Operation *orig, OpBuilder &builder, + MGradientUtils *gutils) const { + auto reduce = cast(orig); + auto &body = reduce.getCombineOp().front(); + + // Get the inner combiner operation (skip terminator) + if (body.getOperations().size() != 2) { + return orig->emitError() + << "Triton reduce autodiff only supports single-op combiners"; + } + Operation &innerOp = body.front(); + + // For max/min reductions, use control flow handler + if (isa(innerOp)) { + llvm::SmallDenseSet operandPositionsToShadow; + llvm::SmallDenseSet resultPositionsToShadow; + for (unsigned i = 0; i < orig->getNumOperands(); i++) { + if (!gutils->isConstantValue(orig->getOperand(i))) + operandPositionsToShadow.insert(i); + } + for (unsigned i = 0; i < orig->getNumResults(); i++) { + if (!gutils->isConstantValue(orig->getResult(i))) + resultPositionsToShadow.insert(i); + } + return mlir::enzyme::detail::controlFlowForwardHandler( + orig, builder, gutils, operandPositionsToShadow, + resultPositionsToShadow); + } + + // Only support sum (addf/addi) reductions + if (!isa(innerOp)) { + return orig->emitError() + << "Triton reduce autodiff only supports add/max/min combiners, " + "got: " + << innerOp.getName(); + } + + // For sum: d(sum(x)) = sum(dx) + // Get the primal operation + Operation *primal = gutils->getNewFromOriginal(orig); + + // Clone the reduce with shadow inputs + IRMapping map; + for (Value input : reduce.getSrcs()) { + if (gutils->isConstantValue(input)) { + // Create zero shadow for constant inputs + auto iface = dyn_cast(input.getType()); + if (!iface) { + return orig->emitError() + << "Cannot differentiate reduce with non-differentiable type"; + } + map.map(input, iface.createNullValue(builder, input.getLoc())); + } else { + map.map(input, gutils->invertPointerM(input, builder)); + } + } + + // Clone entire op with remapped inputs (this correctly handles the region) + Operation *shadowReduce = builder.clone(*orig, map); + + // Clean up: Enzyme may have added shadow block arguments to the primal. + // For sum reduction, we don't need shadows in the combiner body. + // Remove inverted pointers for the inner op result and block arguments. + auto invInnerResult = gutils->invertedPointers.lookup(innerOp.getResult(0)); + gutils->invertedPointers.erase(innerOp.getResult(0)); + if (invInnerResult && invInnerResult.getDefiningOp()) + gutils->erase(invInnerResult.getDefiningOp()); + + // Remove shadow block arguments from primal + auto primalReduce = cast(primal); + auto &primalBody = primalReduce.getCombineOp().front(); + BitVector baToErase(primalBody.getNumArguments()); + for (auto ba : body.getArguments()) { + auto invBA = gutils->invertedPointers.lookup(ba); + if (invBA) { + gutils->invertedPointers.erase(ba); + if (auto invBABlock = dyn_cast(invBA)) { + baToErase.set(invBABlock.getArgNumber()); + } + } + } + primalBody.eraseArguments(baToErase); + + // Set derivatives for results + for (auto [origRes, shadowRes] : + llvm::zip(orig->getResults(), shadowReduce->getResults())) { + if (!gutils->isConstantValue(origRes)) { + gutils->setDiffe(origRes, shadowRes, builder); + } + } + + gutils->eraseIfUnused(orig); + return success(); + } +}; + +/// Forward mode autodiff for Triton DotOp. +/// +/// tt.dot computes: result = a @ b + c +/// Forward mode: d(result) = da @ b + a @ db + dc +/// +/// We chain dot operations using the accumulator: +/// shadow = dc (or zero if c is constant) +/// if a is active: shadow = dot(da, b, shadow) +/// if b is active: shadow = dot(a, db, shadow) +class AutoDiffTritonDotFwd + : public AutoDiffOpInterface::ExternalModel { +public: + LogicalResult createForwardModeTangent(Operation *orig, OpBuilder &builder, + MGradientUtils *gutils) const { + auto dot = cast(orig); + Value a = dot.getA(); + Value b = dot.getB(); + Value c = dot.getC(); + + Value newA = gutils->getNewFromOriginal(a); + Value newB = gutils->getNewFromOriginal(b); + + Value shadow; + if (!gutils->isConstantValue(c)) { + shadow = gutils->invertPointerM(c, builder); + } else { + auto iface = cast(orig->getResult(0).getType()); + shadow = iface.createNullValue(builder, orig->getLoc()); + } + + if (!gutils->isConstantValue(a)) { + Value shadowA = gutils->invertPointerM(a, builder); + IRMapping map; + map.map(a, shadowA); + map.map(b, newB); + map.map(c, shadow); + shadow = builder.clone(*orig, map)->getResult(0); + } + + if (!gutils->isConstantValue(b)) { + Value shadowB = gutils->invertPointerM(b, builder); + IRMapping map; + map.map(a, newA); + map.map(b, shadowB); + map.map(c, shadow); + shadow = builder.clone(*orig, map)->getResult(0); + } + + Value origResult = orig->getResult(0); + if (!gutils->isConstantValue(origResult)) { + gutils->setDiffe(origResult, shadow, builder); + } + + gutils->eraseIfUnused(orig); + return success(); + } +}; + +/// Control flow interface for Triton ReduceOp (needed for max/min) +class AutoDiffTritonReduceCF + : public ControlFlowAutoDiffOpInterface::ExternalModel< + AutoDiffTritonReduceCF, triton::ReduceOp> { +public: + Operation *createWithShadows(Operation *op, OpBuilder &builder, + MGradientUtils *gutils, Operation *original, + ValueRange remappedOperands, + TypeRange rettys) const { + auto reduce = cast(original); + auto newOp = builder.create( + original->getLoc(), remappedOperands, reduce.getAxis()); + IRMapping mapping; + reduce.getCombineOp().cloneInto(&newOp.getCombineOp(), mapping); + return newOp; + } +}; + } // end anonymous namespace void mlir::enzyme::registerTritonDialectAutoDiffInterface( @@ -95,5 +294,8 @@ void mlir::enzyme::registerTritonDialectAutoDiffInterface( triton::FuncOp::attachInterface( *context); triton::PointerType::attachInterface(*context); + triton::DotOp::attachInterface(*context); + triton::ReduceOp::attachInterface(*context); + triton::ReduceOp::attachInterface(*context); }); } diff --git a/src/enzyme_ad/jax/Implementations/TritonDerivatives.td b/src/enzyme_ad/jax/Implementations/TritonDerivatives.td index b51c5a2ba5..7a9258d669 100644 --- a/src/enzyme_ad/jax/Implementations/TritonDerivatives.td +++ b/src/enzyme_ad/jax/Implementations/TritonDerivatives.td @@ -21,10 +21,14 @@ def MakeRange : TritonInst<"MakeRangeOp">; def : TritonReturnOp<"ReturnOp">; def : TritonInactiveOp<"AssertOp">; +def : TritonInactiveOp<"GetProgramIdOp">; def : TritonInactiveOp<"MakeRangeOp">; def : TritonInactiveOp<"PrintOp">; def : ReadOnlyIdentityOp<"triton", "AddPtrOp", [0]>; +def : ReadOnlyIdentityOp<"triton", "AdvanceOp", [0]>; +def : ReadOnlyIdentityOp<"triton", "BroadcastOp", [0]>; +def : ReadOnlyIdentityOp<"triton", "ExpandDimsOp", [0]>; def : ReadOnlyIdentityOp<"triton", "LoadOp", [0]>; def : ReadOnlyIdentityOp<"triton", "SplatOp", [0]>; def : MemoryIdentityOp<"triton", "StoreOp", [1], [0]>; @@ -38,3 +42,9 @@ def : TritonDerivative<"FpToFpOp", (Op $x), (FpToFp (TypeOf $x), (DiffeRet), (FpToFpRoundingMode)) ] >; + +// Note: ReduceOp is handled via custom C++ AutoDiffOpInterface in +// TritonAutoDiffOpInterfaceImpl.cpp because Triton's ReduceOp doesn't +// implement RegionBranchOpInterface required by ControlFlowOp. + +def : RegionTerminatorOp<"triton", "ReduceReturnOp">; diff --git a/test/lit_tests/diffrules/triton/matmul.mlir b/test/lit_tests/diffrules/triton/matmul.mlir new file mode 100644 index 0000000000..b4d9d3abad --- /dev/null +++ b/test/lit_tests/diffrules/triton/matmul.mlir @@ -0,0 +1,115 @@ +// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=matmul_kernel outfn= argTys=enzyme_dup,enzyme_dup,enzyme_dup,enzyme_const,enzyme_const,enzyme_const,enzyme_const,enzyme_const,enzyme_const retTys= mode=ForwardMode" --canonicalize | FileCheck %s + +module { + tt.func public @matmul_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, + %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, + %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32) attributes {noinline = false} { + %c31_i32 = arith.constant 31 : i32 + %c63_i32 = arith.constant 63 : i32 + %c1_i32 = arith.constant 1 : i32 + %cst = arith.constant dense<32> : tensor<64x32xi32> + %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x32xf32> + %cst_1 = arith.constant dense<0.000000e+00> : tensor<64x32xf32> + %c32_i32 = arith.constant 32 : i32 + %c64_i32 = arith.constant 64 : i32 + %true = arith.constant true + %c0_i32 = arith.constant 0 : i32 + %c8_i32 = arith.constant 8 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.addi %arg3, %c63_i32 : i32 + %2 = arith.divsi %1, %c64_i32 : i32 + %3 = arith.addi %arg4, %c31_i32 : i32 + %4 = arith.divsi %3, %c32_i32 : i32 + %5 = arith.muli %4, %c8_i32 : i32 + %6 = arith.divsi %0, %5 : i32 + %7 = arith.muli %6, %c8_i32 : i32 + %8 = arith.subi %2, %7 : i32 + %9 = arith.minsi %8, %c8_i32 : i32 + %10 = arith.remsi %0, %5 : i32 + %11 = arith.remsi %10, %9 : i32 + %12 = arith.addi %7, %11 : i32 + %13 = arith.divsi %10, %9 : i32 + %14 = arith.cmpi sge, %12, %c0_i32 : i32 + llvm.intr.assume %14 : i1 + %15 = arith.cmpi sge, %13, %c0_i32 : i32 + llvm.intr.assume %15 : i1 + %16 = arith.cmpi sgt, %arg6, %c0_i32 : i32 + llvm.intr.assume %16 : i1 + llvm.intr.assume %true : i1 + llvm.intr.assume %true : i1 + %17 = arith.cmpi sgt, %arg7, %c0_i32 : i32 + llvm.intr.assume %17 : i1 + %18 = arith.cmpi sgt, %arg8, %c0_i32 : i32 + llvm.intr.assume %18 : i1 + llvm.intr.assume %true : i1 + %19 = arith.muli %12, %c64_i32 : i32 + %20 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + %21 = tt.splat %19 : i32 -> tensor<64xi32> + %22 = arith.addi %21, %20 : tensor<64xi32> + %23 = tt.splat %arg3 : i32 -> tensor<64xi32> + %24 = arith.remsi %22, %23 : tensor<64xi32> + %25 = arith.muli %13, %c32_i32 : i32 + %26 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32> + %27 = tt.splat %25 : i32 -> tensor<32xi32> + %28 = arith.addi %27, %26 : tensor<32xi32> + %29 = tt.splat %arg4 : i32 -> tensor<32xi32> + %30 = arith.remsi %28, %29 : tensor<32xi32> + %31 = tt.expand_dims %24 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32> + %32 = tt.splat %arg6 : i32 -> tensor<64x1xi32> + %33 = arith.muli %31, %32 : tensor<64x1xi32> + %34 = tt.expand_dims %26 {axis = 0 : i32} : tensor<32xi32> -> tensor<1x32xi32> + %35 = tt.broadcast %33 : tensor<64x1xi32> -> tensor<64x32xi32> + %36 = tt.broadcast %34 : tensor<1x32xi32> -> tensor<64x32xi32> + %37 = arith.addi %35, %36 : tensor<64x32xi32> + %38 = tt.splat %arg0 : !tt.ptr -> tensor<64x32x!tt.ptr> + %39 = tt.addptr %38, %37 : tensor<64x32x!tt.ptr>, tensor<64x32xi32> + %40 = tt.expand_dims %26 {axis = 1 : i32} : tensor<32xi32> -> tensor<32x1xi32> + %41 = tt.splat %arg7 : i32 -> tensor<32x1xi32> + %42 = arith.muli %40, %41 : tensor<32x1xi32> + %43 = tt.expand_dims %30 {axis = 0 : i32} : tensor<32xi32> -> tensor<1x32xi32> + %44 = tt.broadcast %42 : tensor<32x1xi32> -> tensor<32x32xi32> + %45 = tt.broadcast %43 : tensor<1x32xi32> -> tensor<32x32xi32> + %46 = arith.addi %44, %45 : tensor<32x32xi32> + %47 = tt.splat %arg1 : !tt.ptr -> tensor<32x32x!tt.ptr> + %48 = tt.addptr %47, %46 : tensor<32x32x!tt.ptr>, tensor<32x32xi32> + %49 = arith.addi %arg5, %c31_i32 : i32 + %50 = arith.divsi %49, %c32_i32 : i32 + %51:3 = scf.for %arg9 = %c0_i32 to %50 step %c1_i32 iter_args(%arg10 = %cst_1, %arg11 = %39, %arg12 = %48) -> (tensor<64x32xf32>, tensor<64x32x!tt.ptr>, tensor<32x32x!tt.ptr>) : i32 { + %69 = arith.muli %arg9, %c32_i32 : i32 + %70 = arith.subi %arg5, %69 : i32 + %71 = tt.splat %70 : i32 -> tensor<1x32xi32> + %72 = arith.cmpi slt, %34, %71 : tensor<1x32xi32> + %73 = tt.broadcast %72 : tensor<1x32xi1> -> tensor<64x32xi1> + %74 = tt.load %arg11, %73, %cst_1 : tensor<64x32x!tt.ptr> + %75 = tt.splat %70 : i32 -> tensor<32x1xi32> + %76 = arith.cmpi slt, %40, %75 : tensor<32x1xi32> + %77 = tt.broadcast %76 : tensor<32x1xi1> -> tensor<32x32xi1> + %78 = tt.load %arg12, %77, %cst_0 : tensor<32x32x!tt.ptr> + %79 = tt.dot %74, %78, %arg10, inputPrecision = tf32 : tensor<64x32xf32> * tensor<32x32xf32> -> tensor<64x32xf32> + %80 = tt.addptr %arg11, %cst : tensor<64x32x!tt.ptr>, tensor<64x32xi32> + %81 = arith.muli %arg7, %c32_i32 : i32 + %82 = tt.splat %81 : i32 -> tensor<32x32xi32> + %83 = tt.addptr %arg12, %82 : tensor<32x32x!tt.ptr>, tensor<32x32xi32> + scf.yield %79, %80, %83 : tensor<64x32xf32>, tensor<64x32x!tt.ptr>, tensor<32x32x!tt.ptr> + } + %52 = arith.truncf %51#0 : tensor<64x32xf32> to tensor<64x32xf16> + %53 = tt.expand_dims %22 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32> + %54 = tt.splat %arg8 : i32 -> tensor<64x1xi32> + %55 = arith.muli %54, %53 : tensor<64x1xi32> + %56 = tt.splat %arg2 : !tt.ptr -> tensor<64x1x!tt.ptr> + %57 = tt.addptr %56, %55 : tensor<64x1x!tt.ptr>, tensor<64x1xi32> + %58 = tt.expand_dims %28 {axis = 0 : i32} : tensor<32xi32> -> tensor<1x32xi32> + %59 = tt.broadcast %57 : tensor<64x1x!tt.ptr> -> tensor<64x32x!tt.ptr> + %60 = tt.broadcast %58 : tensor<1x32xi32> -> tensor<64x32xi32> + %61 = tt.addptr %59, %60 : tensor<64x32x!tt.ptr>, tensor<64x32xi32> + %62 = tt.splat %arg3 : i32 -> tensor<64x1xi32> + %63 = arith.cmpi slt, %53, %62 : tensor<64x1xi32> + %64 = tt.splat %arg4 : i32 -> tensor<1x32xi32> + %65 = arith.cmpi slt, %58, %64 : tensor<1x32xi32> + %66 = tt.broadcast %63 : tensor<64x1xi1> -> tensor<64x32xi1> + %67 = tt.broadcast %65 : tensor<1x32xi1> -> tensor<64x32xi1> + %68 = arith.andi %66, %67 : tensor<64x32xi1> + tt.store %61, %52, %68 : tensor<64x32x!tt.ptr> + tt.return + } +} diff --git a/test/lit_tests/diffrules/triton/op_coverage.mlir b/test/lit_tests/diffrules/triton/op_coverage.mlir new file mode 100644 index 0000000000..409aa21247 --- /dev/null +++ b/test/lit_tests/diffrules/triton/op_coverage.mlir @@ -0,0 +1,85 @@ +// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=test_triton_ops outfn= argTys=enzyme_dup,enzyme_dup,enzyme_dup retTys=enzyme_dup,enzyme_dup mode=ForwardMode" --canonicalize | FileCheck %s + +module { + tt.func @test_triton_ops(%ptr_in: !tt.ptr, %ptr_out: !tt.ptr, %vec: tensor<64xf32>) -> (tensor<64x32xf32>, f32) { + %pid = tt.get_program_id x : i32 + %range = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> + + // splat + addptr + load + %in_splat = tt.splat %ptr_in : !tt.ptr -> tensor<64x!tt.ptr> + %in_ptrs = tt.addptr %in_splat, %range : tensor<64x!tt.ptr>, tensor<64xi32> + %loaded = tt.load %in_ptrs : tensor<64x!tt.ptr> + + // store loaded data to output pointer + %out_splat = tt.splat %ptr_out : !tt.ptr -> tensor<64x!tt.ptr> + %out_ptrs = tt.addptr %out_splat, %range : tensor<64x!tt.ptr>, tensor<64xi32> + tt.store %out_ptrs, %loaded : tensor<64x!tt.ptr> + + // expand_dims + broadcast to build a matrix from %vec + %col = tt.expand_dims %vec {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> + %mat = tt.broadcast %col : tensor<64x1xf32> -> tensor<64x32xf32> + + // dot: mat @ constant_matrix + zero + %cst = arith.constant dense<1.0> : tensor<32x32xf32> + %zero = arith.constant dense<0.0> : tensor<64x32xf32> + %dot = tt.dot %mat, %cst, %zero, inputPrecision = tf32 : tensor<64x32xf32> * tensor<32x32xf32> -> tensor<64x32xf32> + + // reduce sum over %vec + %sum = "tt.reduce"(%vec) <{axis = 0 : i32}> ({ + ^bb0(%lhs: f32, %rhs: f32): + %add = arith.addf %lhs, %rhs : f32 + tt.reduce.return %add : f32 + }) : (tensor<64xf32>) -> f32 + + tt.return %dot, %sum : tensor<64x32xf32>, f32 + } +} + +// After forward-mode AD with enzyme_dup on all 3 args: +// ptr_in (%arg0) + shadow (%arg1) +// ptr_out (%arg2) + shadow (%arg3) +// vec (%arg4) + shadow vec (%arg5) +// Returns: (primal_dot, shadow_dot, primal_sum, shadow_sum) +// +// The differentiated function interleaves shadow and primal operations: +// - get_program_id is inactive and unused (DCE'd) +// - make_range is inactive (shared by shadow and primal) +// - splat/addptr/load are duplicated for shadow and primal pointers +// - store writes shadow data to shadow ptr, primal data to primal ptr +// - expand_dims/broadcast/dot are duplicated for shadow and primal vec +// - reduce is duplicated for shadow and primal vec + +// CHECK: tt.func @test_triton_ops(%arg0: !tt.ptr, %arg1: !tt.ptr, %arg2: !tt.ptr, %arg3: !tt.ptr, %arg4: tensor<64xf32>, %arg5: tensor<64xf32>) -> (tensor<64x32xf32>, tensor<64x32xf32>, f32, f32) { +// CHECK-NEXT: %[[cst:.+]] = arith.constant dense<0.000000e+00> : tensor<64x32xf32> +// CHECK-NEXT: %[[cst_0:.+]] = arith.constant dense<1.000000e+00> : tensor<32x32xf32> +// CHECK-NEXT: %[[v0:.+]] = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> +// CHECK-NEXT: %[[v1:.+]] = tt.splat %arg1 : !tt.ptr -> tensor<64x!tt.ptr> +// CHECK-NEXT: %[[v2:.+]] = tt.splat %arg0 : !tt.ptr -> tensor<64x!tt.ptr> +// CHECK-NEXT: %[[v3:.+]] = tt.addptr %[[v1]], %[[v0]] : tensor<64x!tt.ptr>, tensor<64xi32> +// CHECK-NEXT: %[[v4:.+]] = tt.addptr %[[v2]], %[[v0]] : tensor<64x!tt.ptr>, tensor<64xi32> +// CHECK-NEXT: %[[v5:.+]] = tt.load %[[v3]] : tensor<64x!tt.ptr> +// CHECK-NEXT: %[[v6:.+]] = tt.load %[[v4]] : tensor<64x!tt.ptr> +// CHECK-NEXT: %[[v7:.+]] = tt.splat %arg3 : !tt.ptr -> tensor<64x!tt.ptr> +// CHECK-NEXT: %[[v8:.+]] = tt.splat %arg2 : !tt.ptr -> tensor<64x!tt.ptr> +// CHECK-NEXT: %[[v9:.+]] = tt.addptr %[[v7]], %[[v0]] : tensor<64x!tt.ptr>, tensor<64xi32> +// CHECK-NEXT: %[[v10:.+]] = tt.addptr %[[v8]], %[[v0]] : tensor<64x!tt.ptr>, tensor<64xi32> +// CHECK-NEXT: tt.store %[[v9]], %[[v5]] : tensor<64x!tt.ptr> +// CHECK-NEXT: tt.store %[[v10]], %[[v6]] : tensor<64x!tt.ptr> +// CHECK-NEXT: %[[v11:.+]] = tt.expand_dims %arg5 {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> +// CHECK-NEXT: %[[v12:.+]] = tt.expand_dims %arg4 {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> +// CHECK-NEXT: %[[v13:.+]] = tt.broadcast %[[v11]] : tensor<64x1xf32> -> tensor<64x32xf32> +// CHECK-NEXT: %[[v14:.+]] = tt.broadcast %[[v12]] : tensor<64x1xf32> -> tensor<64x32xf32> +// CHECK-NEXT: %[[v15:.+]] = tt.dot %[[v13]], %[[cst_0]], %[[cst]], inputPrecision = tf32 : tensor<64x32xf32> * tensor<32x32xf32> -> tensor<64x32xf32> +// CHECK-NEXT: %[[v16:.+]] = tt.dot %[[v14]], %[[cst_0]], %[[cst]], inputPrecision = tf32 : tensor<64x32xf32> * tensor<32x32xf32> -> tensor<64x32xf32> +// CHECK-NEXT: %[[v17:.+]] = "tt.reduce"(%arg5) <{axis = 0 : i32}> ({ +// CHECK-NEXT: ^bb0(%arg6: f32, %arg7: f32): +// CHECK-NEXT: %[[v19:.+]] = arith.addf %arg6, %arg7 : f32 +// CHECK-NEXT: tt.reduce.return %[[v19]] : f32 +// CHECK-NEXT: }) : (tensor<64xf32>) -> f32 +// CHECK-NEXT: %[[v18:.+]] = "tt.reduce"(%arg4) <{axis = 0 : i32}> ({ +// CHECK-NEXT: ^bb0(%arg6: f32, %arg7: f32): +// CHECK-NEXT: %[[v20:.+]] = arith.addf %arg6, %arg7 : f32 +// CHECK-NEXT: tt.reduce.return %[[v20]] : f32 +// CHECK-NEXT: }) : (tensor<64xf32>) -> f32 +// CHECK-NEXT: tt.return %[[v16]], %[[v15]], %[[v18]], %[[v17]] : tensor<64x32xf32>, tensor<64x32xf32>, f32, f32 +// CHECK-NEXT: } diff --git a/test/lit_tests/diffrules/triton/softmax.mlir b/test/lit_tests/diffrules/triton/softmax.mlir index 692ffa9416..3e5b836643 100644 --- a/test/lit_tests/diffrules/triton/softmax.mlir +++ b/test/lit_tests/diffrules/triton/softmax.mlir @@ -1,9 +1,15 @@ -// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=softmax_kernel outfn= argTys=enzyme_dup,enzyme_dup,enzyme_const,enzyme_const,enzyme_const,enzyme_const retTys=enzyme_dup mode=ForwardMode" --canonicalize | FileCheck %s +// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= argTys=enzyme_dup,enzyme_dup,enzyme_const,enzyme_const,enzyme_const,enzyme_const retTys=enzyme_dup,enzyme_dup mode=ForwardMode" --canonicalize | FileCheck %s +// Currently failing within scf.for. module { + // Checking softmax kernel. enzymexla_tt_ext.module @softmax_kernel_tt { builtin.module @softmax_kernel_inner { - tt.func public @softmax_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32, %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + tt.func public @softmax_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, + %arg1: !tt.ptr {tt.divisibility = 16 : i32}, + %arg2: i32 {tt.divisibility = 16 : i32}, + %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32, + %arg5: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { %cst = arith.constant dense<0xFF800000> : tensor<64xf32> %0 = tt.get_program_id x : i32 %1 = tt.get_num_programs x : i32 @@ -41,50 +47,18 @@ module { } } } - func.func @main(%arg0: tensor<1024xf32>, %arg1: tensor<1024xf32>, %arg2: tensor<1024xf32>, %arg3: tensor) -> (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>) { + func.func @main(%arg0: tensor<1024xf32>, + %arg1: tensor<1024xf32>, + %arg2: tensor, + %arg3: tensor, + %arg4: tensor, + %arg5: tensor) -> (tensor<1024xf32>, tensor<1024xf32>) { %c_0 = stablehlo.constant dense<1> : tensor %c_1 = stablehlo.constant dense<16> : tensor - %0:3 = enzymexla_tt_ext.call @softmax_kernel_tt::@softmax_kernel_inner::@softmax_kernel clusters in (%c_0, %c_0, %c_0) blocks in(%c_1, %c_0, %c_0) (%arg0, %arg1, %arg3) {output_operand_aliases = [#stablehlo.output_operand_alias, #stablehlo.output_operand_alias]} : (tensor<1024xf32>, tensor<1024xf32>, tensor) -> (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>) - return %0#0, %0#1, %0#2 : tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32> + %0:2 = enzymexla_tt_ext.call @softmax_kernel_tt::@softmax_kernel_inner::@softmax_kernel clusters in (%c_0, %c_0, %c_0) blocks in(%c_1, %c_0, %c_0) (%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) + {output_operand_aliases = [#stablehlo.output_operand_alias, #stablehlo.output_operand_alias]} : (tensor<1024xf32>, tensor<1024xf32>, tensor, tensor, tensor, tensor) -> (tensor<1024xf32>, tensor<1024xf32>) + return %0#0, %0#1 : tensor<1024xf32>, tensor<1024xf32> } -} - - -// CHECK: tt.func public @fwddiffsoftmax_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: i32 {tt.divisibility = 16 : i32}, %arg3: !tt.ptr {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: !tt.ptr {tt.divisibility = 16 : i32}, %arg6: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { -// CHECK-NEXT: %cst = arith.constant dense<0xFF800000> : tensor<64xf32> loc(#loc1) -// CHECK-NEXT: %0 = tt.get_program_id x : i32 loc(#loc2) -// CHECK-NEXT: %[[v0:.+]] = tt.get_program_id x : i32 loc(#loc2) -// CHECK-NEXT: %1 = tt.get_num_programs x : i32 loc(#loc3) -// CHECK-NEXT: %[[v1:.+]] = tt.get_num_programs x : i32 loc(#loc3) -// CHECK-NEXT: scf.for %arg6 = %[[v0:.+]] to %arg4 step %[[v1:.+]] : i32 { -// CHECK-NEXT: %2 = arith.muli %arg6, %arg2 : i32 loc(#loc5) -// CHECK-NEXT: %3 = tt.addptr %arg1, %2 : !tt.ptr, i32 loc(#loc6) -// CHECK-NEXT: %4 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32> loc(#loc7) -// CHECK-NEXT: %5 = tt.splat %3 : !tt.ptr -> tensor<64x!tt.ptr> loc(#loc8) -// CHECK-NEXT: %6 = tt.addptr %5, %4 : tensor<64x!tt.ptr>, tensor<64xi32> loc(#loc8) -// CHECK-NEXT: %7 = tt.splat %arg5 : i32 -> tensor<64xi32> loc(#loc9) -// CHECK-NEXT: %8 = arith.cmpi slt, %4, %7 : tensor<64xi32> loc(#loc9) -// CHECK-NEXT: %9 = tt.load %6, %8, %cst : tensor<64x!tt.ptr> loc(#loc10) -// CHECK-NEXT: %10 = "tt.reduce"(%9) <{axis = 0 : i32}> ({ -// CHECK-NEXT: ^bb0(%arg7: f32 loc(callsite(#loc1 at #loc12)), %arg8: f32 loc(callsite(#loc1 at #loc12))): -// CHECK-NEXT: %21 = arith.maxnumf %arg7, %arg8 : f32 loc(#loc29) -// CHECK-NEXT: tt.reduce.return %21 : f32 loc(#loc25) -// CHECK-NEXT: }) : (tensor<64xf32>) -> f32 loc(#loc25) -// CHECK-NEXT: %11 = tt.splat %10 : f32 -> tensor<64xf32> loc(#loc14) -// CHECK-NEXT: %12 = arith.subf %9, %11 : tensor<64xf32> loc(#loc14) -// CHECK-NEXT: %13 = math.exp %12 : tensor<64xf32> loc(#loc15) -// CHECK-NEXT: %14 = "tt.reduce"(%13) <{axis = 0 : i32}> ({ -// CHECK-NEXT: ^bb0(%arg7: f32 loc(callsite(#loc1 at #loc17)), %arg8: f32 loc(callsite(#loc1 at #loc17))): -// CHECK-NEXT: %21 = arith.addf %arg7, %arg8 : f32 loc(#loc30) -// CHECK-NEXT: tt.reduce.return %21 : f32 loc(#loc27) -// CHECK-NEXT: }) : (tensor<64xf32>) -> f32 loc(#loc27) -// CHECK-NEXT: %15 = tt.splat %14 : f32 -> tensor<64xf32> loc(#loc19) -// CHECK-NEXT: %16 = arith.divf %13, %15 : tensor<64xf32> loc(#loc19) -// CHECK-NEXT: %17 = arith.muli %arg6, %arg3 : i32 loc(#loc20) -// CHECK-NEXT: %18 = tt.addptr %arg0, %17 : !tt.ptr, i32 loc(#loc21) -// CHECK-NEXT: %19 = tt.splat %18 : !tt.ptr -> tensor<64x!tt.ptr> loc(#loc22) -// CHECK-NEXT: %20 = tt.addptr %19, %4 : tensor<64x!tt.ptr>, tensor<64xi32> loc(#loc22) -// CHECK-NEXT: tt.store %20, %16, %8 : tensor<64x!tt.ptr> loc(#loc23) -// CHECK-NEXT: } {tt.num_stages = 2 : i32} loc(#loc4) -// CHECK-NEXT: tt.return loc(#loc24) -// CHECK-NEXT: } +} \ No newline at end of file diff --git a/test/lit_tests/triton/matmul_kernel.mlir b/test/lit_tests/triton/matmul_kernel.mlir new file mode 100644 index 0000000000..8472a09ce2 --- /dev/null +++ b/test/lit_tests/triton/matmul_kernel.mlir @@ -0,0 +1,37 @@ +// RUN: enzymexlamlir-opt %s -canonicalize | FileCheck %s + +module { + // CHECK: enzymexla_tt_ext.module + enzymexla_tt_ext.module @add_kernel_tt { + builtin.module @add_kernel_inner { + tt.func public @add_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %c1024_i32 = arith.constant 1024 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c1024_i32 : i32 + %2 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %3 = tt.splat %1 : i32 -> tensor<1024xi32> + %4 = arith.addi %3, %2 : tensor<1024xi32> + %5 = tt.splat %arg3 : i32 -> tensor<1024xi32> + %6 = arith.cmpi slt, %4, %5 : tensor<1024xi32> + %7 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %8 = tt.addptr %7, %4 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %9 = tt.load %8, %6 : tensor<1024x!tt.ptr> + %10 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %11 = tt.addptr %10, %4 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %12 = tt.load %11, %6 : tensor<1024x!tt.ptr> + %13 = arith.addf %9, %12 : tensor<1024xf32> + %14 = tt.splat %arg2 : !tt.ptr -> tensor<1024x!tt.ptr> + %15 = tt.addptr %14, %4 : tensor<1024x!tt.ptr>, tensor<1024xi32> + tt.store %15, %13, %6 : tensor<1024x!tt.ptr> + tt.return + } + } + } + func.func @main(%arg0: tensor<1024xf32>, %arg1: tensor<1024xf32>, %arg2: tensor<1024xf32>, %arg3: tensor) -> (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>) { + %c_0 = stablehlo.constant dense<1> : tensor + %c_1 = stablehlo.constant dense<16> : tensor + // CHECK: enzymexla_tt_ext.call + %0:3 = enzymexla_tt_ext.call @add_kernel_tt::@add_kernel_inner::@add_kernel clusters in (%c_0, %c_0, %c_0) blocks in(%c_1, %c_0, %c_0) (%arg0, %arg1, %arg2, %arg3) {output_operand_aliases = [#stablehlo.output_operand_alias, #stablehlo.output_operand_alias, #stablehlo.output_operand_alias]} : (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor) -> (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>) + return %0#0, %0#1, %0#2 : tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32> + } +} diff --git a/test/lit_tests/triton/softmax_kernel.mlir b/test/lit_tests/triton/softmax_kernel.mlir new file mode 100644 index 0000000000..8472a09ce2 --- /dev/null +++ b/test/lit_tests/triton/softmax_kernel.mlir @@ -0,0 +1,37 @@ +// RUN: enzymexlamlir-opt %s -canonicalize | FileCheck %s + +module { + // CHECK: enzymexla_tt_ext.module + enzymexla_tt_ext.module @add_kernel_tt { + builtin.module @add_kernel_inner { + tt.func public @add_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %c1024_i32 = arith.constant 1024 : i32 + %0 = tt.get_program_id x : i32 + %1 = arith.muli %0, %c1024_i32 : i32 + %2 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> + %3 = tt.splat %1 : i32 -> tensor<1024xi32> + %4 = arith.addi %3, %2 : tensor<1024xi32> + %5 = tt.splat %arg3 : i32 -> tensor<1024xi32> + %6 = arith.cmpi slt, %4, %5 : tensor<1024xi32> + %7 = tt.splat %arg0 : !tt.ptr -> tensor<1024x!tt.ptr> + %8 = tt.addptr %7, %4 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %9 = tt.load %8, %6 : tensor<1024x!tt.ptr> + %10 = tt.splat %arg1 : !tt.ptr -> tensor<1024x!tt.ptr> + %11 = tt.addptr %10, %4 : tensor<1024x!tt.ptr>, tensor<1024xi32> + %12 = tt.load %11, %6 : tensor<1024x!tt.ptr> + %13 = arith.addf %9, %12 : tensor<1024xf32> + %14 = tt.splat %arg2 : !tt.ptr -> tensor<1024x!tt.ptr> + %15 = tt.addptr %14, %4 : tensor<1024x!tt.ptr>, tensor<1024xi32> + tt.store %15, %13, %6 : tensor<1024x!tt.ptr> + tt.return + } + } + } + func.func @main(%arg0: tensor<1024xf32>, %arg1: tensor<1024xf32>, %arg2: tensor<1024xf32>, %arg3: tensor) -> (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>) { + %c_0 = stablehlo.constant dense<1> : tensor + %c_1 = stablehlo.constant dense<16> : tensor + // CHECK: enzymexla_tt_ext.call + %0:3 = enzymexla_tt_ext.call @add_kernel_tt::@add_kernel_inner::@add_kernel clusters in (%c_0, %c_0, %c_0) blocks in(%c_1, %c_0, %c_0) (%arg0, %arg1, %arg2, %arg3) {output_operand_aliases = [#stablehlo.output_operand_alias, #stablehlo.output_operand_alias, #stablehlo.output_operand_alias]} : (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor) -> (tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32>) + return %0#0, %0#1, %0#2 : tensor<1024xf32>, tensor<1024xf32>, tensor<1024xf32> + } +}