diff --git a/compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp b/compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp
index 3eaa7b3..66cfa40 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp
@@ -22,6 +22,8 @@
 
 namespace mlir::iree_compiler {
 
+#define CEILDIV(a, b) ((a + b - 1) / b)
+
 #define GEN_PASS_DEF_TILEANDDISTRIBUTETOWORKGROUPSUSINGFORALLOPPASS
 #include "iree/compiler/Codegen/Common/Passes.h.inc"
 
@@ -119,6 +121,140 @@
   return llvm::to_vector(llvm::reverse(mapping));
 }
 
+//===---------------------------------------------------------------------===//
+// Post tiling cleanup patterns
+//===---------------------------------------------------------------------===//
+
+/// Prune the values corresponding to the dropped loops.
+static SmallVector<OpFoldResult>
+pruneDroppedLoops(ArrayRef<OpFoldResult> inputs,
+                  const llvm::SmallDenseSet<int> &droppedLoops) {
+  SmallVector<OpFoldResult> prunedInputs;
+  for (auto [index, input] : llvm::enumerate(inputs)) {
+    if (droppedLoops.contains(index)) {
+      continue;
+    }
+    prunedInputs.push_back(input);
+  }
+  return prunedInputs;
+}
+
+/// Prune the mapping attributes corresponding to the dropped loops.
+/// Note that we cant just drop them. We need to rebalance the
+/// attributes so that the workgroup attributes are perfectly ordered.
+/// For example, if the attribute list is
+///
+/// ```
+/// [workgroup_mapping<x>, workgroup_mapping<z:1>,
+///  workgroup_mapping<z>, workgroup_mapping<y>,
+///  workgroup_mapping<z:3>, workgroup_mapping<z:2>]
+/// ```
+///
+/// and the droppedloops are `{1, 3}`, then the new mapping should be
+///
+/// ```
+/// [workgroup_mapping<x>, workgroup_mapping<y>,
+///  workgroup_mapping<z:1>, workgroup_mapping<z>]
+/// ```
+SmallVector<Attribute>
+pruneDroppedLoops(ArrayRef<Attribute> inputs,
+                  const llvm::SmallDenseSet<int> &droppedLoops) {
+  SmallVector<IREE::Codegen::WorkgroupMappingAttr> droppedMappings;
+  SmallVector<Attribute> prunedAttrs;
+  for (auto [index, input] : llvm::enumerate(inputs)) {
+    if (droppedLoops.contains(index)) {
+      droppedMappings.push_back(
+          cast<IREE::Codegen::WorkgroupMappingAttr>(input));
+    } else {
+      prunedAttrs.push_back(input);
+    }
+  }
+  for (auto droppedMapping : droppedMappings) {
+    for (auto [index, prunedAttr] : llvm::enumerate(prunedAttrs)) {
+      auto prunedMappingAttr =
+          cast<IREE::Codegen::WorkgroupMappingAttr>(prunedAttr);
+      if (droppedMapping < prunedMappingAttr) {
+        prunedAttrs[index] =
+            IREE::Codegen::WorkgroupMappingAttr::getAttributeFromMappingId(
+                prunedAttr.getContext(), prunedMappingAttr.getMappingId() - 1);
+      }
+    }
+  }
+  return prunedAttrs;
+}
+
+/// Find dimensions of the loop that are unit-trip count and drop them from the
+/// distributed dimensions.
+static LogicalResult dropUnitDistributedDims(RewriterBase &rewriter,
+                                             scf::ForallOp forallOp) {
+  SmallVector<OpFoldResult> mixedLbs = forallOp.getMixedLowerBound();
+  SmallVector<OpFoldResult> mixedUbs = forallOp.getMixedUpperBound();
+  SmallVector<OpFoldResult> mixedSteps = forallOp.getMixedStep();
+
+  // Find the index of loops to be dropped.
+  llvm::SmallDenseSet<int> droppedLoops;
+  for (auto [index, lb, ub, step] :
+       llvm::enumerate(mixedLbs, mixedUbs, mixedSteps)) {
+    if (!isa<Attribute>(lb) || !isa<Attribute>(ub) || !isa<Attribute>(step)) {
+      continue;
+    }
+    int64_t lbVal = getConstantIntValue(lb).value();
+    int64_t ubVal = getConstantIntValue(ub).value();
+    int64_t stepVal = getConstantIntValue(step).value();
+    if (CEILDIV(ubVal - lbVal, stepVal) == 1) {
+      droppedLoops.insert(index);
+    }
+  }
+  if (droppedLoops.empty()) {
+    return success();
+  }
+
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(forallOp);
+  SmallVector<OpFoldResult> newLbs =
+      pruneDroppedLoops(ArrayRef<OpFoldResult>(mixedLbs), droppedLoops);
+  SmallVector<OpFoldResult> newUbs =
+      pruneDroppedLoops(ArrayRef<OpFoldResult>(mixedUbs), droppedLoops);
+  SmallVector<OpFoldResult> newSteps =
+      pruneDroppedLoops(ArrayRef<OpFoldResult>(mixedSteps), droppedLoops);
+  std::optional<ArrayAttr> newMapping;
+  if (auto currMapping = forallOp.getMapping()) {
+    SmallVector<Attribute> newMappingAttrs =
+        pruneDroppedLoops(currMapping.value().getValue(), droppedLoops);
+    newMapping = rewriter.getArrayAttr(newMappingAttrs);
+  }
+
+  Value zero = rewriter.create<arith::ConstantIndexOp>(forallOp.getLoc(), 0);
+  auto newForallOp = rewriter.create<scf::ForallOp>(
+      forallOp.getLoc(), newLbs, newUbs, newSteps, forallOp.getInits(),
+      newMapping, [](OpBuilder &, Location, ValueRange) {});
+
+  SmallVector<Value> argReplacements;
+  int newLoopBlockArgNum = 0;
+  auto newLoopBodyArgs = newForallOp.getInductionVars();
+  for (auto [index, oldBlockArg] :
+       llvm::enumerate(forallOp.getInductionVars())) {
+    if (droppedLoops.contains(index)) {
+      argReplacements.push_back(zero);
+    } else {
+      argReplacements.push_back(newLoopBodyArgs[newLoopBlockArgNum++]);
+    }
+  }
+  argReplacements.append(newForallOp.getRegionIterArgs().begin(),
+                         newForallOp.getRegionIterArgs().end());
+
+  Block *oldLoopBody = forallOp.getBody();
+  Block *newLoopBody = newForallOp.getBody();
+  rewriter.mergeBlocks(oldLoopBody, newLoopBody, argReplacements);
+
+  rewriter.replaceOp(forallOp, newForallOp.getResults());
+  return success();
+}
+
+//===---------------------------------------------------------------------===//
+// Pass implementation.
+//===---------------------------------------------------------------------===//
+
 void TileAndDistributeToWorkgroupsUsingForallOpPass::runOnOperation() {
   auto funcOp = getOperation();
   auto *context = &getContext();
@@ -180,6 +316,19 @@
     }
     std::swap(tileAndFuseResult->loops, tilingLoops);
   }
+  if (!tilingLoops.empty()) {
+    if (tilingLoops.size() != 1 || !isa<scf::ForallOp>(tilingLoops[0])) {
+      funcOp.emitOpError(
+          "expected tiling to produce a single `scf.forall` loop");
+      return signalPassFailure();
+    }
+
+    auto forallOp = cast<scf::ForallOp>(tilingLoops[0]);
+    if (failed(dropUnitDistributedDims(rewriter, forallOp))) {
+      forallOp.emitOpError("failed to drop unit dimensions");
+      return signalPassFailure();
+    }
+  }
 
   // Cleanup patterns for tile and distribute
   {
@@ -191,6 +340,7 @@
     context->getOrLoadDialect<IREE::LinalgExt::IREELinalgExtDialect>()
         ->getCanonicalizationPatterns(patterns);
     memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
+    scf::ForallOp::getCanonicalizationPatterns(patterns, context);
     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns)))) {
       funcOp.emitOpError("tiling canonicalization failed");
       return signalPassFailure();
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_workgroups_using_forall.mlir b/compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_workgroups_using_forall.mlir
index cc6d728..fb68210 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_workgroups_using_forall.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_workgroups_using_forall.mlir
@@ -135,22 +135,6 @@
 
 // -----
 
-func.func @gemm_unit_N(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x1xf32>,
-    %arg2 : tensor<?x1xf32>) -> tensor<?x1xf32> {
-  %0 = linalg.matmul {
-      lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64, 64, 64]]>}
-      ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x1xf32>)
-      outs(%arg2 : tensor<?x1xf32>) -> tensor<?x1xf32>
-  return %0 : tensor<?x1xf32>
-}
-// CHECK-LABEL: func.func @gemm_unit_N(
-//       CHECK:   %[[RESULT:.+]] = scf.forall (%[[IV0:[a-zA-Z0-9]+]])
-//       CHECK:     %[[MATMUL:.+]] = linalg.matmul
-//       CHECK:     scf.forall.in_parallel {
-//       CHECK:       tensor.parallel_insert_slice %[[MATMUL]] into %{{.+}}[%[[IV0]], 0] [%{{.+}}, 1]
-
-// -----
-
 func.func @gemm_unit_M_unit_N(%arg0 : tensor<1x1xf32>, %arg1 : tensor<1x1xf32>,
     %arg2 : tensor<1x1xf32>) -> tensor<1x1xf32> {
   %0 = linalg.matmul {
@@ -351,3 +335,74 @@
 //       CHECK:     scf.forall.in_parallel
 //       CHECK:       tensor.parallel_insert_slice %[[MATMUL]]
 //       CHECK:   return %[[RESULT]]
+
+// -----
+
+func.func @avoid_unit_range_distribute(
+    %arg0 : tensor<32x?x?x16x16xf16>, %arg1 : tensor<32x?x8x16x16xf16>,
+    %arg2 : tensor<32x?x16x8x16xf16>) -> tensor<32x?x16x8x16xf16> {
+   %0 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d5, d2, d6)>,
+                       affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d5, d3, d6, d4)>,
+                       affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4)>],
+      iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
+      ins(%arg0, %arg1 : tensor<32x?x?x16x16xf16>, tensor<32x?x8x16x16xf16>)
+      outs(%arg2 : tensor<32x?x16x8x16xf16>)
+      attrs =  {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[1, 1, 16, 1, 16, 1, 16]]>} {
+    ^bb0(%b0: f16, %b1: f16, %b2 : f16):
+      %1 = arith.mulf %b0, %b1 : f16
+      %2 = arith.addf %b2, %1 : f16
+      linalg.yield %2 : f16
+  } -> tensor<32x?x16x8x16xf16>
+  return %0 : tensor<32x?x16x8x16xf16>
+}
+// CHECK-LABEL: func @avoid_unit_range_distribute(
+//       CHECK:   scf.forall (%{{[a-zA-Z0-9]+}}, %{{[a-zA-Z0-9]+}}, %{{[a-zA-Z0-9]+}}) in (32, %{{.+}}, 8)
+//       CHECK:   mapping = [#iree_codegen.workgroup_mapping<z>, #iree_codegen.workgroup_mapping<y>, #iree_codegen.workgroup_mapping<x>]
+
+// -----
+
+// This just verifies that constant dim propagation works as expected after tiling.
+func.func @set_size_to_tilesize_when_divisible(
+    %arg0 : tensor<?x16x32x128xf16>, %arg1 : tensor<4096x32x28xf16>,
+    %arg2 : tensor<?x16x4096xf16>) -> tensor<?x16x4096xf16> {
+   %0 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3, d4)>,
+                       affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>,
+                       affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>],
+      iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]}
+      ins(%arg0, %arg1 : tensor<?x16x32x128xf16>, tensor<4096x32x28xf16>)
+      outs(%arg2 : tensor<?x16x4096xf16>)
+      attrs =  {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[1, 16, 128, 1, 128]]>} {
+    ^bb0(%b0: f16, %b1: f16, %b2 : f16):
+      %1 = arith.mulf %b0, %b1 : f16
+      %2 = arith.addf %b2, %1 : f16
+      linalg.yield %2 : f16
+  } -> tensor<?x16x4096xf16>
+  return %0 : tensor<?x16x4096xf16>
+}
+// CHECK-LABEL: func @set_size_to_tilesize_when_divisible(
+//       CHECK:   scf.forall
+//       CHECK:     %[[GENERIC:.+]] = linalg.generic
+//       CHECK:     scf.forall.in_parallel
+//       CHECK:       tensor.parallel_insert_slice %[[GENERIC]]
+//  CHECK-SAME:           tensor<1x16x128xf16> into tensor<?x16x4096xf16>
+
+// -----
+
+// This just verifies that constant dim propagation works as expected after tiling.
+func.func @generate_no_distribution(%arg0 : tensor<16xf16>) -> tensor<16xf16> {
+   %empty = tensor.empty() : tensor<16xf16>
+   %0 = linalg.generic {
+      indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+      iterator_types = ["parallel"]}
+      ins(%arg0 : tensor<16xf16>) outs(%empty : tensor<16xf16>)
+      attrs =  {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[16]]>} {
+    ^bb0(%b0: f16, %b1: f16):
+      %1 = arith.mulf %b0, %b0 : f16
+      linalg.yield %1 : f16
+  } -> tensor<16xf16>
+  return %0 : tensor<16xf16>
+}
+// CHECK-LABEL: func @generate_no_distribution(
+//   CHECK-NOT:   scf.forall
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp
index b818ac8..5f4be34 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp
@@ -382,6 +382,19 @@
   return success();
 }
 
+WorkgroupMappingAttr
+WorkgroupMappingAttr::getAttributeFromMappingId(MLIRContext *context,
+                                                int64_t mappingId) {
+  int64_t linearizedDim = std::max<int64_t>(mappingId - 2, 0);
+  WorkgroupId id =
+      symbolizeWorkgroupId(std::min<uint64_t>(mappingId, 2)).value();
+  return WorkgroupMappingAttr::get(context, id, linearizedDim);
+}
+
+bool WorkgroupMappingAttr::operator<(const WorkgroupMappingAttr &rhs) const {
+  return getMappingId() < rhs.getMappingId();
+}
+
 LogicalResult WorkgroupMappingAttr::verifyAttrList(
     MLIRContext *context, function_ref<InFlightDiagnostic()> emitError,
     ArrayRef<Attribute> attrs) {
@@ -404,14 +417,8 @@
     mappingAttrs.push_back(typedAttr);
   }
 
-  llvm::sort(mappingAttrs, [](const IREE::Codegen::WorkgroupMappingAttr &lhs,
-                              const IREE::Codegen::WorkgroupMappingAttr &rhs) {
-    if (lhs.getId() != rhs.getId()) {
-      return lhs.getId() < rhs.getId();
-    }
-    assert(lhs.getId() == IREE::Codegen::WorkgroupId::IdZ);
-    return lhs.getDelinearizedDim() < rhs.getDelinearizedDim();
-  });
+  llvm::sort(mappingAttrs);
+
   // First element has to be `workgroup_mapping<x>`.
   if (mappingAttrs.front().getId() != IREE::Codegen::WorkgroupId::IdX) {
     return emitError() << "missing `workgroup_mapping<x>`";
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td
index 95ea6e5..6190b6c 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td
@@ -190,6 +190,15 @@
     static LogicalResult verifyAttrList(::mlir::MLIRContext *context,
       ::llvm::function_ref<::mlir::InFlightDiagnostic ()> emitError,
       ArrayRef<Attribute> attrs);
+
+    // Convert from mapping ID to attribute.
+    static ::mlir::iree_compiler::IREE::Codegen::WorkgroupMappingAttr
+    getAttributeFromMappingId(MLIRContext *context, int64_t mappingId);
+
+    // Less than operator for easy comparison.
+    bool operator<(
+        const ::mlir::iree_compiler::IREE::Codegen::WorkgroupMappingAttr &rhs)
+        const;
   }];
 
   let genVerifyDecl = 1;