diff --git a/src/enzyme_ad/jax/Passes/AffineCFG.cpp b/src/enzyme_ad/jax/Passes/AffineCFG.cpp index 3ba9f7f1b2..9f702b06c9 100644 --- a/src/enzyme_ad/jax/Passes/AffineCFG.cpp +++ b/src/enzyme_ad/jax/Passes/AffineCFG.cpp @@ -612,6 +612,9 @@ AffineApplyNormalizer::AffineApplyNormalizer(AffineMap map, continue; } if (auto idx = decast.getDefiningOp()) { + // Sign extension of i1 flips the value (true -> -1). + if (idx.getIn().getType().isInteger(1)) + break; decast = idx.getIn(); continue; } @@ -1447,7 +1450,8 @@ bool isValidIndex(Value val, Region *scope) { return isValidIndex(cast.getOperand(), scope); if (auto cast = val.getDefiningOp()) - return isValidIndex(cast.getOperand(), scope); + if (!cast.getOperand().getType().isInteger(1)) + return isValidIndex(cast.getOperand(), scope); if (auto cast = val.getDefiningOp()) return isValidIndex(cast.getOperand(), scope); diff --git a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp index e3346bd36e..41826d9539 100644 --- a/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp +++ b/src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp @@ -3442,6 +3442,10 @@ tryRaisingOpToStableHLO(Operation *op, IRMapping &mapping, OpBuilder &builder, return success(); } + // An optimizer hint carries no semantics a tensor program needs. + if (isa(op)) + return success(); + return op->emitError("cannot raise op to stablehlo") << *op; } @@ -3592,6 +3596,273 @@ struct AffineToStableHLORaisingPass AffineToStableHLORaisingPass> { using AffineToStableHLORaisingBase::AffineToStableHLORaisingBase; + // A branch yielding one of several buffers blocks raising: no tensor can + // stand for "one of these two memrefs". Duplicate the branch at every + // access instead — a load becomes a value-yielding branch loading in each + // arm, a store becomes a store in each arm — so every access reaches a + // real buffer and the usual select/mask raising applies. Branch bodies are + // cloned per access, so only effect-free bodies qualify. + static void expandBufferBranches(Operation *root) { + // The same shape also arrives as an arith.select of two buffers: expand + // each access into an scf.if on the select's condition. + SmallVector selects; + root->walk([&](arith::SelectOp sel) { + if (isa(sel.getType())) + selects.push_back(sel); + }); + for (auto sel : selects) { + for (OpOperand &use : llvm::make_early_inc_range(sel->getUses())) { + Operation *user = use.getOwner(); + bool isLoad = isa(user); + bool isStore = isa(user); + unsigned memIdx = isLoad ? 0 : 1; + if ((!isLoad && !isStore) || use.getOperandNumber() != memIdx) + continue; + OpBuilder b(user); + auto newIf = scf::IfOp::create( + b, user->getLoc(), + isLoad ? TypeRange(user->getResult(0).getType()) : TypeRange(), + sel.getCondition(), /*withElseRegion=*/true); + auto fillArm = [&](Value buf, Block *dstArm) { + dstArm->clear(); + IRMapping m; + OpBuilder ab = OpBuilder::atBlockBegin(dstArm); + Operation *access = ab.clone(*user, m); + access->setOperand(memIdx, buf); + scf::YieldOp::create(ab, user->getLoc(), + isLoad ? ValueRange(access->getResult(0)) + : ValueRange()); + }; + fillArm(sel.getTrueValue(), newIf.thenBlock()); + fillArm(sel.getFalseValue(), newIf.elseBlock()); + if (isLoad) + user->getResult(0).replaceAllUsesWith(newIf.getResult(0)); + user->erase(); + } + if (sel->use_empty()) + sel.erase(); + } + + // The coefficient ternary arrives as an scf.if yielding a pointer whose + // arms compute constant- or index-offset geps; push each viewing access + // down into a clone of the branch so no pointer crosses the yield. + SmallVector scfWorklist; + root->walk([&](scf::IfOp ifOp) { + if (ifOp.elseBlock() && llvm::any_of(ifOp.getResultTypes(), [](Type t) { + return isa(t); + })) + scfWorklist.push_back(ifOp); + }); + // Arms may be cloned once per pushed-down access, so they must not + // write; reads are idempotent and safe to duplicate. + auto armClonable = [](Block *b) { + return llvm::all_of(b->without_terminator(), [](Operation &op) { + if (isMemoryEffectFree(&op)) + return true; + auto mem = dyn_cast(&op); + if (!mem || op.getNumRegions() != 0) + return false; + SmallVector effects; + mem.getEffects(effects); + return llvm::all_of(effects, [](MemoryEffects::EffectInstance &e) { + return isa(e.getEffect()); + }); + }); + }; + for (auto ifOp : scfWorklist) { + Block *thenB = ifOp.thenBlock(), *elseB = ifOp.elseBlock(); + if (!armClonable(thenB) || !armClonable(elseB)) + continue; + for (auto [i, res] : llvm::enumerate(ifOp.getResults())) { + if (!isa(res.getType())) + continue; + Value thenV = thenB->getTerminator()->getOperand(i); + Value elseV = elseB->getTerminator()->getOperand(i); + for (OpOperand &use : llvm::make_early_inc_range(res.getUses())) { + auto p2m = dyn_cast(use.getOwner()); + if (!p2m) + continue; + for (OpOperand &ause : llvm::make_early_inc_range(p2m->getUses())) { + Operation *acc = ause.getOwner(); + bool isLoad = isa(acc); + bool isStore = isa(acc); + unsigned memIdx = isLoad ? 0 : 1; + if ((!isLoad && !isStore) || ause.getOperandNumber() != memIdx) + continue; + OpBuilder b(acc); + auto newIf = scf::IfOp::create( + b, acc->getLoc(), + isLoad ? TypeRange(acc->getResult(0).getType()) : TypeRange(), + ifOp.getCondition(), /*withElseRegion=*/true); + auto fillArm = [&](Block *srcArm, Value yielded, Block *dstArm) { + dstArm->clear(); + IRMapping m; + OpBuilder ab = OpBuilder::atBlockBegin(dstArm); + for (Operation &armOp : srcArm->without_terminator()) + ab.clone(armOp, m); + Operation *view = ab.clone(*p2m.getOperation(), m); + view->setOperand(0, m.lookupOrDefault(yielded)); + Operation *access = ab.clone(*acc, m); + access->setOperand(memIdx, view->getResult(0)); + scf::YieldOp::create(ab, acc->getLoc(), + isLoad ? ValueRange(access->getResult(0)) + : ValueRange()); + }; + fillArm(thenB, thenV, newIf.thenBlock()); + fillArm(elseB, elseV, newIf.elseBlock()); + if (isLoad) + acc->getResult(0).replaceAllUsesWith(newIf.getResult(0)); + acc->erase(); + } + if (p2m->use_empty()) + p2m.erase(); + } + } + // Rebuild without dead pointer results if scalars keep it alive. + if (llvm::all_of(ifOp.getResults(), + [](Value r) { return r.use_empty(); })) { + ifOp.erase(); + continue; + } + if (llvm::any_of(ifOp.getResults(), [](Value r) { + return isa(r.getType()) && r.use_empty(); + })) { + SmallVector liveIdx; + SmallVector liveTypes; + for (auto [i, res] : llvm::enumerate(ifOp.getResults())) { + if (isa(res.getType()) && res.use_empty()) + continue; + liveIdx.push_back((unsigned)i); + liveTypes.push_back(res.getType()); + } + OpBuilder b(ifOp); + auto newIf = + scf::IfOp::create(b, ifOp.getLoc(), liveTypes, ifOp.getCondition(), + /*withElseRegion=*/true); + auto rebuildArm = [&](Block *srcArm, Block *dstArm) { + dstArm->clear(); + IRMapping m; + OpBuilder ab = OpBuilder::atBlockBegin(dstArm); + for (Operation &armOp : srcArm->without_terminator()) + ab.clone(armOp, m); + SmallVector yields; + for (unsigned i : liveIdx) + yields.push_back( + m.lookupOrDefault(srcArm->getTerminator()->getOperand(i))); + scf::YieldOp::create(ab, ifOp.getLoc(), yields); + for (Operation &armOp : + llvm::make_early_inc_range(dstArm->without_terminator())) + if (armOp.use_empty() && isMemoryEffectFree(&armOp)) + armOp.erase(); + }; + rebuildArm(thenB, newIf.thenBlock()); + rebuildArm(elseB, newIf.elseBlock()); + for (auto [k, i] : llvm::enumerate(liveIdx)) + ifOp.getResult(i).replaceAllUsesWith(newIf.getResult(k)); + ifOp.erase(); + } + } + + SmallVector worklist; + root->walk([&](affine::AffineIfOp ifOp) { + if (ifOp.hasElse() && llvm::any_of(ifOp.getResultTypes(), [](Type t) { + return isa(t); + })) + worklist.push_back(ifOp); + }); + for (auto ifOp : worklist) { + Block *thenB = ifOp.getThenBlock(), *elseB = ifOp.getElseBlock(); + if (!armClonable(thenB) || !armClonable(elseB)) + continue; + for (auto [i, res] : llvm::enumerate(ifOp.getResults())) { + if (!isa(res.getType())) + continue; + Value thenV = thenB->getTerminator()->getOperand(i); + Value elseV = elseB->getTerminator()->getOperand(i); + for (OpOperand &use : llvm::make_early_inc_range(res.getUses())) { + Operation *user = use.getOwner(); + bool isLoad = isa(user); + bool isStore = isa(user); + unsigned memIdx = isLoad ? 0 : 1; + if ((!isLoad && !isStore) || use.getOperandNumber() != memIdx) + continue; + OpBuilder b(user); + auto newIf = affine::AffineIfOp::create( + b, user->getLoc(), + isLoad ? TypeRange(user->getResult(0).getType()) : TypeRange(), + ifOp.getIntegerSet(), ifOp.getOperands(), + /*withElseRegion=*/true); + auto fillArm = [&](Block *srcArm, Value yielded, Block *dstArm) { + if (Operation *term = dstArm->empty() ? nullptr : &dstArm->back()) + if (term->hasTrait()) + term->erase(); + IRMapping m; + OpBuilder ab = OpBuilder::atBlockEnd(dstArm); + for (Operation &armOp : srcArm->without_terminator()) + ab.clone(armOp, m); + Operation *access = ab.clone(*user, m); + access->setOperand(memIdx, m.lookupOrDefault(yielded)); + affine::AffineYieldOp::create( + ab, user->getLoc(), + isLoad ? ValueRange(access->getResult(0)) : ValueRange()); + }; + fillArm(thenB, thenV, newIf.getThenBlock()); + fillArm(elseB, elseV, newIf.getElseBlock()); + if (isLoad) + user->getResult(0).replaceAllUsesWith(newIf.getResult(0)); + user->erase(); + } + } + if (llvm::all_of(ifOp.getResults(), + [](Value r) { return r.use_empty(); })) { + ifOp.erase(); + continue; + } + // Scalar results may keep the branch alive; rebuild it without the + // now-dead buffer results so no unraisable cast lingers in the arms. + if (llvm::any_of(ifOp.getResults(), [](Value r) { + return isa(r.getType()) && r.use_empty(); + })) { + SmallVector liveIdx; + SmallVector liveTypes; + for (auto [i, res] : llvm::enumerate(ifOp.getResults())) { + if (isa(res.getType()) && res.use_empty()) + continue; + liveIdx.push_back((unsigned)i); + liveTypes.push_back(res.getType()); + } + OpBuilder b(ifOp); + auto newIf = affine::AffineIfOp::create( + b, ifOp.getLoc(), liveTypes, ifOp.getIntegerSet(), + ifOp.getOperands(), /*withElseRegion=*/true); + auto rebuildArm = [&](Block *srcArm, Block *dstArm) { + if (Operation *term = dstArm->empty() ? nullptr : &dstArm->back()) + if (term->hasTrait()) + term->erase(); + IRMapping m; + OpBuilder ab = OpBuilder::atBlockEnd(dstArm); + for (Operation &armOp : srcArm->without_terminator()) + ab.clone(armOp, m); + SmallVector yields; + for (unsigned i : liveIdx) + yields.push_back( + m.lookupOrDefault(srcArm->getTerminator()->getOperand(i))); + affine::AffineYieldOp::create(ab, ifOp.getLoc(), yields); + // The buffer arms may still hold the dead casts; drop them. + for (Operation &armOp : + llvm::make_early_inc_range(dstArm->without_terminator())) + if (armOp.use_empty() && isMemoryEffectFree(&armOp)) + armOp.erase(); + }; + rebuildArm(thenB, newIf.getThenBlock()); + rebuildArm(elseB, newIf.getElseBlock()); + for (auto [k, i] : llvm::enumerate(liveIdx)) + ifOp.getResult(i).replaceAllUsesWith(newIf.getResult(k)); + ifOp.erase(); + } + } + } + // An access does not care about the address space of its base, but the // raising identifies buffers by SSA root: a memory_space_cast view would // split one buffer into two roots and lose store propagation. Retarget the @@ -3900,6 +4171,7 @@ struct AffineToStableHLORaisingPass // actually raises. for (auto func : funcs) { stripAccessMemorySpaceCasts(func); + expandBufferBranches(func); boundParallelAxes(func); peelDynamicParallelDims(func); } @@ -3923,6 +4195,7 @@ struct AffineToStableHLORaisingPass op->walk([&](enzymexla::GPUWrapperOp g) { gwrap.push_back(g); }); for (auto g : gwrap) { stripAccessMemorySpaceCasts(g); + expandBufferBranches(g); boundParallelAxes(g); peelDynamicParallelDims(g); } diff --git a/src/enzyme_ad/jax/Passes/ArithRaising.cpp b/src/enzyme_ad/jax/Passes/ArithRaising.cpp index 77b6da1f00..62c7616fa3 100644 --- a/src/enzyme_ad/jax/Passes/ArithRaising.cpp +++ b/src/enzyme_ad/jax/Passes/ArithRaising.cpp @@ -110,6 +110,18 @@ struct RaiseToConvert : public OpRewritePattern { if (!ty) return failure(); + // stablehlo.convert reads i1 as boolean (true -> 1), but a sign + // extension of i1 means true -> -1: negate the boolean's conversion. + if (std::is_same_v && + cast(op.getIn().getType()) + .getElementType() + .isInteger(1)) { + Value conv = + stablehlo::ConvertOp::create(rewriter, op.getLoc(), ty, op.getIn()); + rewriter.replaceOpWithNewOp(op, conv); + return success(); + } + rewriter.replaceOpWithNewOp(op, ty, op.getIn()); return success(); } @@ -682,6 +694,7 @@ struct ArithRaisingPass RaiseUnary, RaiseUnary, RaiseUnary, + RaiseUnary, RaiseUnary, RaiseUnary, RaiseUnary, diff --git a/src/enzyme_ad/jax/Passes/AutoBatching.cpp b/src/enzyme_ad/jax/Passes/AutoBatching.cpp index 01121bb184..cc9f1321b1 100644 --- a/src/enzyme_ad/jax/Passes/AutoBatching.cpp +++ b/src/enzyme_ad/jax/Passes/AutoBatching.cpp @@ -2086,8 +2086,75 @@ WhileIsCopySimplify::matchAndRewriteImpl(stablehlo::WhileOp whileOp, auto dusInductionVarDims = getInductionVariableDimension(dusOp, affineIndexInfo, whileOp, info); + Value indirectScatterIndices; + if (dusInductionVarDims.empty()) { + // Handle a common non-contiguous copy idiom: + // + // dst[index_table[iv]] = src[iv] + // + // The destination index is not affine, so it cannot be raised as one + // large dynamic_update_slice. If the index table is a constant set of + // unique indices, hoist its per-iteration load and use it to build a + // scatter instead. Requiring uniqueness preserves the sequential + // overwrite semantics of the original loop. + int64_t indirectDim = -1; + stablehlo::DynamicSliceOp indexSlice; + for (auto [dim, startIndex] : llvm::enumerate(dusOp.getStartIndices())) { + if (info.isConstantAcrossIterations(startIndex, false)) + continue; + if (indirectDim != -1) { + indirectDim = -1; + break; + } + + Value indexValue = startIndex; + if (auto reshape = indexValue.getDefiningOp()) + indexValue = reshape.getOperand(); + indexSlice = indexValue.getDefiningOp(); + if (!indexSlice) + break; + indirectDim = dim; + } + + DenseElementsAttr indexTable; + if (indirectDim >= 0 && indexSlice && + matchPattern(indexSlice.getOperand(), m_Constant(&indexTable)) && + indexTable.getNumElements() == + static_cast(info.getConstantNumIters())) { + llvm::SmallDenseSet seenIndices; + bool allUniqueAndInBounds = true; + int64_t maxIndex = cast(dusOp.getOperand().getType()) + .getDimSize(indirectDim) - + cast(dusOp.getUpdate().getType()) + .getDimSize(indirectDim); + for (APInt index : indexTable.getValues()) { + if (index.isNegative() || index.getSExtValue() > maxIndex || + !seenIndices.insert(index).second) { + allUniqueAndInBounds = false; + break; + } + } + + auto indexInductionDims = getInductionVariableDimension( + indexSlice, affineIndexInfo, whileOp, info); + if (allUniqueAndInBounds && indexInductionDims.size() == 1 && + info.canHoistOperationFromLoop(indexSlice, indexInductionDims)) { + rewriter.setInsertionPoint(whileOp); + if (info.hoistOperationFromLoop(rewriter, indexSlice.getOperand(), + indexSlice, indexInductionDims, + indirectScatterIndices)) { + indirectScatterIndices = stablehlo::ReshapeOpCreate( + rewriter, whileOp.getLoc(), indirectScatterIndices, + {info.getConstantNumIters(), 1}); + dusInductionVarDims.push_back(indirectDim); + } + } + } + } + if (dusInductionVarDims.empty() || - !info.canHoistOperationFromLoop(dusOp, dusInductionVarDims)) { + (!indirectScatterIndices && + !info.canHoistOperationFromLoop(dusOp, dusInductionVarDims))) { continue; } @@ -2234,11 +2301,36 @@ WhileIsCopySimplify::matchAndRewriteImpl(stablehlo::WhileOp whileOp, } Value newDUS; - bool successfulHoist = info.hoistOperationFromLoop( - rewriter, whileOp.getOperands()[idx], newDUSUpdate, dusOp, - dusInductionVarDims, newDUS); - if (!successfulHoist) { - return failure(); + if (indirectScatterIndices) { + auto operandTy = + cast(whileOp.getOperands()[idx].getType()); + auto updateTy = cast(newDUSUpdate.getType()); + SmallVector updateWindowDims(updateTy.getRank() - 1); + std::iota(updateWindowDims.begin(), updateWindowDims.end(), 1); + + auto scatter = stablehlo::ScatterOp::create( + rewriter, dusOp.getLoc(), ValueRange{whileOp.getOperands()[idx]}, + indirectScatterIndices, ValueRange{newDUSUpdate}, + stablehlo::ScatterDimensionNumbersAttr::get( + dusOp.getContext(), updateWindowDims, dusInductionVarDims, + /*inputBatchingDims=*/{}, /*scatterIndicesBatchingDims=*/{}, + dusInductionVarDims, /*indexVectorDim=*/1), + /*indicesAreSorted=*/false, /*uniqueIndices=*/true); + Block *updateBody = rewriter.createBlock( + &scatter.getUpdateComputation(), {}, + {RankedTensorType::get({}, operandTy.getElementType()), + RankedTensorType::get({}, operandTy.getElementType())}, + {dusOp.getLoc(), dusOp.getLoc()}); + rewriter.setInsertionPointToStart(updateBody); + stablehlo::ReturnOp::create(rewriter, dusOp.getLoc(), + updateBody->getArgument(1)); + newDUS = scatter.getResult(0); + } else { + bool successfulHoist = info.hoistOperationFromLoop( + rewriter, whileOp.getOperands()[idx], newDUSUpdate, dusOp, + dusInductionVarDims, newDUS); + if (!successfulHoist) + return failure(); } whileOp.getResult(idx).replaceAllUsesWith(newDUS); diff --git a/src/enzyme_ad/jax/Passes/LibDeviceFuncsRaisingPass.cpp b/src/enzyme_ad/jax/Passes/LibDeviceFuncsRaisingPass.cpp index ff053141c7..27c0da1e5b 100644 --- a/src/enzyme_ad/jax/Passes/LibDeviceFuncsRaisingPass.cpp +++ b/src/enzyme_ad/jax/Passes/LibDeviceFuncsRaisingPass.cpp @@ -700,6 +700,17 @@ using ConvertFMFMathFromLLVMPattern = using AbsFOpLowering = ConvertFMFMathFromLLVMPattern; + +// llvm.intr.abs carries an is_int_min_poison flag arith has no place for; +// drop it and raise to math.absi. +struct AbsIOpRaising : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(LLVM::AbsOp op, + PatternRewriter &rewriter) const override { + rewriter.replaceOpWithNewOp(op, op.getIn()); + return success(); + } +}; using CeilOpLowering = ConvertFMFMathFromLLVMPattern; using CopySignOpLowering = @@ -1415,6 +1426,7 @@ void populateLLVMToMathPatterns(MLIRContext *context, // From // https://github.com/llvm/llvm-project/blob/7d8b4eb0ead277f41ff69525ed807f9f6e227f37/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp#L306 // patterns.add(converter); + patterns.add(patterns.getContext()); patterns.add -1; stablehlo.convert reads i1 as +// boolean (true -> 1), so the raised form negates the conversion. + +// CHECK-LABEL: @extsi_bool +// CHECK: %[[C:.+]] = stablehlo.convert %arg0 : (tensor<4xi1>) -> tensor<4xi32> +// CHECK: %[[N:.+]] = stablehlo.negate %[[C]] : tensor<4xi32> +// CHECK: return %[[N]] + +// CHECK-LABEL: @extsi_wide +// CHECK: stablehlo.convert +// CHECK-NOT: stablehlo.negate + +// CHECK-LABEL: @extui_bool +// CHECK: stablehlo.convert +// CHECK-NOT: stablehlo.negate + +module { + func.func @extsi_bool(%arg0: tensor<4xi1>) -> tensor<4xi32> { + %0 = arith.extsi %arg0 : tensor<4xi1> to tensor<4xi32> + return %0 : tensor<4xi32> + } + func.func @extsi_wide(%arg0: tensor<4xi8>) -> tensor<4xi32> { + %0 = arith.extsi %arg0 : tensor<4xi8> to tensor<4xi32> + return %0 : tensor<4xi32> + } + func.func @extui_bool(%arg0: tensor<4xi1>) -> tensor<4xi32> { + %0 = arith.extui %arg0 : tensor<4xi1> to tensor<4xi32> + return %0 : tensor<4xi32> + } +} diff --git a/test/lit_tests/autobatching/scatter_non_contiguous.mlir b/test/lit_tests/autobatching/scatter_non_contiguous.mlir new file mode 100644 index 0000000000..21bdd2a022 --- /dev/null +++ b/test/lit_tests/autobatching/scatter_non_contiguous.mlir @@ -0,0 +1,72 @@ +// RUN: enzymexlamlir-opt --enzyme-hlo-opt="enable_auto_batching_passes=true" %s | FileCheck %s + +// CHECK-LABEL: func.func @scatter_non_contiguous +// CHECK-NOT: stablehlo.while +// CHECK: %[[INDICES:.*]] = stablehlo.constant dense<{{.*}}> : tensor<4x1xi32> +// CHECK: %[[RESULT:.*]] = "stablehlo.scatter"(%{{.*}}, %[[INDICES]], %arg0) +// CHECK-SAME: inserted_window_dims = [0] +// CHECK-SAME: scatter_dims_to_operand_dims = [0] +// CHECK-SAME: index_vector_dim = 1 +// CHECK-SAME: unique_indices = true +// CHECK: stablehlo.return %{{.*}} : tensor +// CHECK: return %[[RESULT]] : tensor<4xf64> +func.func @scatter_non_contiguous(%arg0: tensor<4xf64>) -> tensor<4xf64> { + %zero_i32 = stablehlo.constant dense<0> : tensor + %zero_i64 = stablehlo.constant dense<0> : tensor + %one_i64 = stablehlo.constant dense<1> : tensor + %four_i64 = stablehlo.constant dense<4> : tensor + %init = stablehlo.constant dense<0.0> : tensor<4xf64> + %permutation = stablehlo.constant dense<[2, 0, 3, 1]> : tensor<4xi32> + %result:2 = stablehlo.while(%iv = %zero_i64, %output = %init) + : tensor, tensor<4xf64> + cond { + %pred = stablehlo.compare LT, %iv, %four_i64 + : (tensor, tensor) -> tensor + stablehlo.return %pred : tensor + } do { + %next = stablehlo.add %iv, %one_i64 : tensor + %source_index = stablehlo.convert %iv : (tensor) -> tensor + %update = stablehlo.dynamic_slice %arg0, %source_index, sizes = [1] + : (tensor<4xf64>, tensor) -> tensor<1xf64> + %index_slice = stablehlo.dynamic_slice %permutation, %iv, sizes = [1] + : (tensor<4xi32>, tensor) -> tensor<1xi32> + %index = stablehlo.reshape %index_slice + : (tensor<1xi32>) -> tensor + %updated = stablehlo.dynamic_update_slice %output, %update, %index + : (tensor<4xf64>, tensor<1xf64>, tensor) -> tensor<4xf64> + stablehlo.return %next, %updated : tensor, tensor<4xf64> + } + return %result#1 : tensor<4xf64> +} + +// Replacing sequential overwrites by an overwrite scatter is only valid when +// indices are unique. Repeated destinations must therefore retain the loop. +// CHECK-LABEL: func.func @scatter_repeated_indices +// CHECK: stablehlo.while +func.func @scatter_repeated_indices(%arg0: tensor<4xf64>) -> tensor<4xf64> { + %zero_i64 = stablehlo.constant dense<0> : tensor + %one_i64 = stablehlo.constant dense<1> : tensor + %four_i64 = stablehlo.constant dense<4> : tensor + %init = stablehlo.constant dense<0.0> : tensor<4xf64> + %indices = stablehlo.constant dense<[2, 0, 2, 1]> : tensor<4xi32> + %result:2 = stablehlo.while(%iv = %zero_i64, %output = %init) + : tensor, tensor<4xf64> + cond { + %pred = stablehlo.compare LT, %iv, %four_i64 + : (tensor, tensor) -> tensor + stablehlo.return %pred : tensor + } do { + %next = stablehlo.add %iv, %one_i64 : tensor + %source_index = stablehlo.convert %iv : (tensor) -> tensor + %update = stablehlo.dynamic_slice %arg0, %source_index, sizes = [1] + : (tensor<4xf64>, tensor) -> tensor<1xf64> + %index_slice = stablehlo.dynamic_slice %indices, %iv, sizes = [1] + : (tensor<4xi32>, tensor) -> tensor<1xi32> + %index = stablehlo.reshape %index_slice + : (tensor<1xi32>) -> tensor + %updated = stablehlo.dynamic_update_slice %output, %update, %index + : (tensor<4xf64>, tensor<1xf64>, tensor) -> tensor<4xf64> + stablehlo.return %next, %updated : tensor, tensor<4xf64> + } + return %result#1 : tensor<4xf64> +} diff --git a/test/lit_tests/mem2reg_distinct_symbols.mlir b/test/lit_tests/mem2reg_distinct_symbols.mlir new file mode 100644 index 0000000000..83ab0e2895 --- /dev/null +++ b/test/lit_tests/mem2reg_distinct_symbols.mlir @@ -0,0 +1,24 @@ +// RUN: enzymexlamlir-opt %s -polygeist-mem2reg -split-input-file | FileCheck %s + +// Two accesses whose affine maps agree land on the same slot only when the +// symbols bound to the map agree as well; a load indexed by a different +// symbol names a different location and must not be forwarded. +func.func @distinct_symbols(%i: index, %s1: index, %s2: index, %v: f64) -> f64 { + %c1 = llvm.mlir.constant(1 : i32) : i32 + %mem = llvm.alloca %c1 x !llvm.array<81 x f64> : (i32) -> !llvm.ptr + %view = "enzymexla.pointer2memref"(%mem) : (!llvm.ptr) -> memref + affine.store %v, %view[%i + symbol(%s1) * 9] : memref + %a = affine.load %view[%i + symbol(%s1) * 9] : memref + %b = affine.load %view[%i + symbol(%s2) * 9] : memref + %r = arith.addf %a, %b : f64 + return %r : f64 +} + +// The store and the first load bind the same symbol, so that load takes the +// stored value; the second load binds another symbol and stays a load. +// CHECK-LABEL: func.func @distinct_symbols( +// CHECK-SAME: %[[I:[a-z0-9]+]]: index, %[[S1:[a-z0-9]+]]: index, %[[S2:[a-z0-9]+]]: index, %[[V:[a-z0-9]+]]: f64 +// CHECK: affine.store %[[V]], %{{.*}}[%[[I]] + symbol(%[[S1]]) * 9] +// CHECK: %[[B:.+]] = affine.load %{{.*}}[%[[I]] + symbol(%[[S2]]) * 9] +// CHECK: %[[R:.+]] = arith.addf %[[V]], %[[B]] : f64 +// CHECK: return %[[R]] : f64 diff --git a/test/lit_tests/raising/absi.mlir b/test/lit_tests/raising/absi.mlir new file mode 100644 index 0000000000..5e05075831 --- /dev/null +++ b/test/lit_tests/raising/absi.mlir @@ -0,0 +1,20 @@ +// RUN: enzymexlamlir-opt --libdevice-funcs-raise %s | FileCheck %s --check-prefix=RAISE +// RUN: enzymexlamlir-opt --arith-raise %s | FileCheck %s --check-prefix=HLO + +module { + // RAISE-LABEL: @intr_absi + // RAISE: math.absi %arg0 : i32 + // RAISE-NOT: llvm.intr.abs + func.func @intr_absi(%arg0: i32) -> i32 { + %res = "llvm.intr.abs"(%arg0) <{is_int_min_poison = false}> : (i32) -> i32 + func.return %res : i32 + } + + // HLO-LABEL: @tensor_absi + // HLO: stablehlo.abs %arg0 : tensor<20xi32> + // HLO-NOT: math.absi + func.func @tensor_absi(%arg0: tensor<20xi32>) -> tensor<20xi32> { + %res = math.absi %arg0 : tensor<20xi32> + func.return %res : tensor<20xi32> + } +} diff --git a/test/lit_tests/raising/buffer_branch_expand.mlir b/test/lit_tests/raising/buffer_branch_expand.mlir new file mode 100644 index 0000000000..0e924eacca --- /dev/null +++ b/test/lit_tests/raising/buffer_branch_expand.mlir @@ -0,0 +1,38 @@ +// RUN: enzymexlamlir-opt %s --raise-affine-to-stablehlo | FileCheck %s + +// A branch yielding whole buffers (mfem's `const_coeff ? c0 : c` ternary) +// cannot raise as a value; expand each access into the branch so only +// scalars cross the yield. + +// CHECK-LABEL: @ifbuf_raised +// CHECK: stablehlo.select + +module { + func.func private @sel(%out: memref<16xf64, 1>, %a: memref<16xf64, 1>, %b: memref<16xf64, 1>, %flag: memref<1xi1, 1>) { + affine.parallel (%t) = (0) to (16) { + %c = affine.load %flag[0] : memref<1xi1, 1> + %sel = arith.select %c, %a, %b : memref<16xf64, 1> + %v = affine.load %sel[%t] : memref<16xf64, 1> + affine.store %v, %out[%t] : memref<16xf64, 1> + } + return + } + + // CHECK-LABEL: @sel_raised + // CHECK: stablehlo.select + // CHECK-NOT: arith.select + func.func private @ifbuf(%out: memref<16xf64, 1>, %a: memref<16xf64, 1>, %b: memref<16xf64, 1>, %n: memref<1xi32, 1>) { + %nv = affine.load %n[0] : memref<1xi32, 1> + %ni = arith.index_cast %nv : i32 to index + affine.parallel (%t) = (0) to (16) { + %buf = affine.if affine_set<()[s0] : (s0 - 1 >= 0)>()[%ni] -> memref<16xf64, 1> { + affine.yield %a : memref<16xf64, 1> + } else { + affine.yield %b : memref<16xf64, 1> + } + %v = affine.load %buf[%t] : memref<16xf64, 1> + affine.store %v, %out[%t] : memref<16xf64, 1> + } + return + } +} diff --git a/test/lit_tests/raising/ignore_assume.mlir b/test/lit_tests/raising/ignore_assume.mlir new file mode 100644 index 0000000000..c2519a0def --- /dev/null +++ b/test/lit_tests/raising/ignore_assume.mlir @@ -0,0 +1,21 @@ +// RUN: enzymexlamlir-opt %s --raise-affine-to-stablehlo | FileCheck %s + +// An optimizer hint carries no semantics a tensor program needs; a kernel +// carrying llvm.intr.assume still raises. + +// CHECK-LABEL: @with_assume +// CHECK-NOT: llvm.intr.assume + +module { + func.func private @with_assume(%out: memref<16xf64, 1>, %in: memref<16xf64, 1>, %nbuf: memref<1xi32, 1>) { + %c0_i32 = arith.constant 0 : i32 + affine.parallel (%t) = (0) to (16) { + %n = affine.load %nbuf[0] : memref<1xi32, 1> + %pos = arith.cmpi sgt, %n, %c0_i32 : i32 + "llvm.intr.assume"(%pos) <{op_bundle_sizes = array, op_bundle_tags = []}> : (i1) -> () + %v = affine.load %in[%t] : memref<16xf64, 1> + affine.store %v, %out[%t] : memref<16xf64, 1> + } + return + } +} diff --git a/test/lit_tests/raising/raise_buffer_select.mlir b/test/lit_tests/raising/raise_buffer_select.mlir index 901901a973..5e9d3b62e7 100644 --- a/test/lit_tests/raising/raise_buffer_select.mlir +++ b/test/lit_tests/raising/raise_buffer_select.mlir @@ -1,7 +1,7 @@ // RUN: enzymexlamlir-opt %s --raise-affine-to-stablehlo --split-input-file | FileCheck %s -// A uniform branch choosing between two read-only buffers raises as a select -// of the whole tensors. +// A uniform branch choosing between two read-only buffers expands per access +// and raises as a select of the gathered values. func.func @bufsel(%a: memref<100xf64, 1>, %b: memref<100xf64, 1>, %out: memref<100xf64, 1>, %flagbuf: memref) { %f = affine.load %flagbuf[] : memref %fi = arith.index_cast %f : i64 to index @@ -18,5 +18,5 @@ func.func @bufsel(%a: memref<100xf64, 1>, %b: memref<100xf64, 1>, %out: memref<1 } // CHECK-LABEL: func.func private @bufsel_raised( -// CHECK: stablehlo.select %{{.+}}, %arg0, %arg1 : tensor, tensor<100xf64> +// CHECK: stablehlo.select %{{.+}}, %{{.+}}, %{{.+}} : tensor<100xi1>, tensor<100xf64> diff --git a/workspace.bzl b/workspace.bzl index 712523334c..ec8c5034fc 100644 --- a/workspace.bzl +++ b/workspace.bzl @@ -1,7 +1,7 @@ JAX_COMMIT = "cec06d116c05f0d52adfceee3d3b730fdbcb0ce5" JAX_SHA256 = "" -ENZYME_COMMIT = "fc6bb335b90ef09c2c16b413a408cafcc836086b" +ENZYME_COMMIT = "277ecb5335a75883852c59436e358f82939dd10d" ENZYME_SHA256 = "" ML_TOOLCHAIN_COMMIT = "30ef4a9096f9490e8f198faa5ce5bbddd1b72fdb"