diff --git a/compiler/src/iree/compiler/Codegen/Common/FissionTransferOpsInControlFlow.cpp b/compiler/src/iree/compiler/Codegen/Common/FissionTransferOpsInControlFlow.cpp
index 48779f8..6c72bd6 100644
--- a/compiler/src/iree/compiler/Codegen/Common/FissionTransferOpsInControlFlow.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/FissionTransferOpsInControlFlow.cpp
@@ -10,6 +10,7 @@
 #include "iree/compiler/Codegen/Utils/Utils.h"
 #include "llvm/Support/DebugLog.h"
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/PatternMatch.h"
@@ -32,11 +33,17 @@
                                      vector::TransferReadOp readOp,
                                      scf::ForOp forOp) {
   auto loc = forOp.getLoc();
-  auto allocaSize = rewriter.create<arith::CeilDivUIOp>(
-      loc,
-      rewriter.create<arith::SubIOp>(loc, forOp.getUpperBound(),
-                                     forOp.getLowerBound()),
-      forOp.getStep());
+  // We use an affine map to calculate the size of the alloca rather
+  // than arith ops directly due to an issue with finding the static upper bound
+  // when using arith.ceildivui described in
+  // https://github.com/iree-org/iree/issues/21872
+  AffineExpr lb, ub, step;
+  bindSymbols(rewriter.getContext(), lb, ub, step);
+  AffineMap sizeMap = AffineMap::get(0, 3, (ub - lb).ceilDiv(step));
+  auto allocaSize = affine::makeComposedFoldedAffineApply(
+      rewriter, loc, sizeMap,
+      getAsOpFoldResult(ValueRange{forOp.getLowerBound(), forOp.getUpperBound(),
+                                   forOp.getStep()}));
 
   auto vectorType = cast<VectorType>(readOp.getVectorType());
   SmallVector<int64_t> memrefShape(vectorType.getShape());
@@ -46,8 +53,9 @@
   auto memrefType = MemRefType::get(memrefShape, vectorType.getElementType(),
                                     AffineMap{}, privateAddrSpaceAttr);
 
-  return rewriter.create<memref::AllocaOp>(loc, memrefType,
-                                           ValueRange{allocaSize});
+  return rewriter.create<memref::AllocaOp>(
+      loc, memrefType,
+      ValueRange{getValueOrCreateConstantIndexOp(rewriter, loc, allocaSize)});
 }
 
 /// Creates an index for accessing the memref in the loop. This index is
diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.td b/compiler/src/iree/compiler/Codegen/Common/Passes.td
index d0d38ec..bc34858 100644
--- a/compiler/src/iree/compiler/Codegen/Common/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Common/Passes.td
@@ -370,7 +370,8 @@
   let summary =
       "Fission transfer read and write ops in control flow to allow prefetching.";
   let dependentDialects = [
-      "memref::MemRefDialect"
+      "memref::MemRefDialect",
+      "affine::AffineDialect"
   ];
   let options = [
     Option<"FissionMultiTrip", "fission-multi-trip",
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/fission_transfer_ops_control_flow.mlir b/compiler/src/iree/compiler/Codegen/Common/test/fission_transfer_ops_control_flow.mlir
index cfca25d..99cae6c 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/fission_transfer_ops_control_flow.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/fission_transfer_ops_control_flow.mlir
@@ -1,5 +1,5 @@
-// RUN: iree-opt --split-input-file -pass-pipeline="builtin.module(func.func(iree-codegen-fission-transfer-ops-in-control-flow{fission-multi-trip}))" %s | FileCheck %s --check-prefixes=CHECK-ALL,MULTI
-// RUN: iree-opt --split-input-file -pass-pipeline="builtin.module(func.func(iree-codegen-fission-transfer-ops-in-control-flow))" %s | FileCheck %s --check-prefixes=CHECK-ALL,SINGLE
+// RUN: iree-opt --split-input-file -pass-pipeline="builtin.module(func.func(iree-codegen-fission-transfer-ops-in-control-flow{fission-multi-trip}))" --mlir-print-local-scope %s | FileCheck %s --check-prefixes=CHECK-ALL,MULTI
+// RUN: iree-opt --split-input-file -pass-pipeline="builtin.module(func.func(iree-codegen-fission-transfer-ops-in-control-flow))" --mlir-print-local-scope %s | FileCheck %s --check-prefixes=CHECK-ALL,SINGLE
 
 // CHECK-ALL-LABEL: @fission_global_read_to_private_write
 // CHECK-ALL-SAME: %[[ARG0:.*]]: memref<1x?x?x8xbf16, #amdgpu.address_space<fat_raw_buffer>>
@@ -66,8 +66,7 @@
   }
   return
 }
-// MULTI: %[[SUB:.*]] = arith.subi %c16, %[[ARG0]]
-// MULTI: %[[DIV:.*]] = arith.ceildivui %[[SUB]], %c128
+// MULTI: %[[DIV:.*]] = affine.apply affine_map<()[s0] -> ((-s0 + 16) ceildiv 128)>()[%[[ARG0]]]
 // MULTI: %[[ALLOCA:.*]] = memref.alloca(%[[DIV]])
 // MULTI: scf.for %[[ITER:.*]] = %[[ARG0]] to %c16 step %c128 {
 // MULTI:   %[[READ:.*]] = vector.transfer_read %arg1[%c0, %c0], %cst {in_bounds = [true, true]}