diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
index d036b9c..7bb6437 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
+++ b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -53,9 +53,11 @@
         "Passes.cpp",
         "RemoveZeroExtentTensors.cpp",
         "SetEncoding.cpp",
+        "Utils.cpp",
     ],
     hdrs = [
         "Passes.h",
+        "Utils.h",
     ],
     deps = [
         ":PassHeaders",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
index 88781e7..ea706e1 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
+++ b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -38,6 +38,7 @@
     GlobalOptimization
   HDRS
     "Passes.h"
+    "Utils.h"
   SRCS
     "Convert1X1FilterConv2DToMatmul.cpp"
     "DetachElementwiseFromNamedOps.cpp"
@@ -48,6 +49,7 @@
     "Passes.cpp"
     "RemoveZeroExtentTensors.cpp"
     "SetEncoding.cpp"
+    "Utils.cpp"
   DEPS
     ::PassHeaders
     ::PassesIncGen
diff --git a/compiler/src/iree/compiler/GlobalOptimization/SetEncoding.cpp b/compiler/src/iree/compiler/GlobalOptimization/SetEncoding.cpp
index 342c19d..789e6a5 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/SetEncoding.cpp
+++ b/compiler/src/iree/compiler/GlobalOptimization/SetEncoding.cpp
@@ -15,10 +15,12 @@
 #include "iree/compiler/Codegen/Dialect/IREECodegenAttrs.h"
 #include "iree/compiler/GlobalOptimization/PassDetail.h"
 #include "iree/compiler/GlobalOptimization/Passes.h"
+#include "iree/compiler/GlobalOptimization/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
@@ -136,17 +138,61 @@
       getAttr(narrow.M), getAttr(narrow.N));
 }
 
-static Value padAndSetEncoding(OpBuilder &builder, Location loc, Value source,
-                               IREE::LinalgExt::EncodingUser user,
-                               IREE::LinalgExt::EncodingRole role,
-                               TypeRange operandTypes,
-                               MatmulNarrowSizes narrow) {
-  // No need to specify original_type in the encoding passed to pad(), because
+// Creates a linalg::GenericOp that performs an element-wise cast of the same
+// type as performed in `castOp`, and returns the result enceoded with
+// `encodingAttr`. The element type of `encoded` is expected to be the same as
+// the element type of the input to `castOp`, which can be a CastOpInterface op
+// on a tensor or single element.
+static Value castEncodedResult(OpBuilder &builder, Location loc, Value encoded,
+                               CastOpInterface castOp,
+                               IREE::LinalgExt::EncodingAttr encodingAttr) {
+  auto encodedType = cast<RankedTensorType>(encoded.getType());
+  auto castResultElemType = getElementTypeOrSelf(castOp->getResultTypes()[0]);
+  auto castedType = RankedTensorType::get(encodedType.getShape(),
+                                          castResultElemType, encodingAttr);
+  assert(encodedType.getElementType() ==
+             getElementTypeOrSelf(castOp->getOperandTypes()[0]) &&
+         "Expected encoded element type to be the same as the cast input "
+         "element type");
+  SmallVector<OpFoldResult> inputMixedSizes =
+      tensor::getMixedSizes(builder, loc, encoded);
+  Value init = builder.create<tensor::EmptyOp>(
+      loc, inputMixedSizes, castResultElemType, encodingAttr);
+  SmallVector<AffineMap> maps(
+      2, AffineMap::getMultiDimIdentityMap(castedType.getRank(),
+                                           builder.getContext()));
+  SmallVector<utils::IteratorType> iteratorTypes(castedType.getRank(),
+                                                 utils::IteratorType::parallel);
+  auto genericOp = castOp->getParentOfType<linalg::GenericOp>();
+  NamedAttrList castAttrs =
+      genericOp ? genericOp->getAttrs() : castOp->getAttrs();
+  return builder
+      .create<linalg::GenericOp>(
+          loc, castedType, encoded, init, maps, iteratorTypes,
+          [&](OpBuilder &b, Location nestedLoc, ValueRange args) {
+            Value castRes =
+                b.create(nestedLoc, castOp->getName().getIdentifier(), args[0],
+                         castResultElemType)
+                    ->getResult(0);
+            b.create<linalg::YieldOp>(nestedLoc, castRes);
+          },
+          castAttrs)
+      ->getResult(0);
+}
+
+static Value
+padAndSetEncoding(OpBuilder &builder, Location loc, Value source,
+                  IREE::LinalgExt::EncodingUser user,
+                  IREE::LinalgExt::EncodingRole role, TypeRange operandTypes,
+                  MatmulNarrowSizes narrow,
+                  std::optional<CastOpInterface> castOp = std::nullopt) {
+  Value padSource = castOp ? source.getDefiningOp()->getOperand(0) : source;
+  // No need to specify original_type in the encoding poadded to pad(), because
   // the operand there is the `source` tensor, so it will default to reading its
   // original shape.
   auto encodingForPad = makeEncoding(builder, user, role, operandTypes,
                                      /*originalType=*/Type{}, narrow);
-  Value padded = pad(builder, loc, source, encodingForPad);
+  Value padded = pad(builder, loc, padSource, encodingForPad);
   // For setEncoding() below, we potentially need to specify an encoding with an
   // explicit original_type, because the operand there is the padded tensor
   // returned by pad() above, but we want setEncoding to be aware of the
@@ -154,11 +200,16 @@
   // verbosity, we only specify the original original_type when it differs from
   // the tensor type that the encoding is applied to.
   auto encodingForSetEncoding = encodingForPad;
-  if (padded.getType() != source.getType()) {
+  if (padded.getType() != padSource.getType()) {
     encodingForSetEncoding = makeEncoding(builder, user, role, operandTypes,
-                                          source.getType(), narrow);
+                                          padSource.getType(), narrow);
   }
-  return setEncoding(builder, loc, padded, encodingForSetEncoding);
+  Value encoded = setEncoding(builder, loc, padded, encodingForSetEncoding);
+  if (castOp) {
+    encoded = castEncodedResult(builder, loc, encoded, castOp.value(),
+                                encodingForSetEncoding);
+  }
+  return encoded;
 }
 
 static Value unsetEncodingAndExtractSlice(OpBuilder &builder, Location loc,
@@ -297,8 +348,12 @@
       }
       return {};
     };
-    Type lhsElemType = getElemType(origLhs);
-    Type rhsElemType = getElemType(origRhs);
+    std::optional<CastOpInterface> maybeLhsCastOp = getDefiningCastOp(origLhs);
+    std::optional<CastOpInterface> maybeRhsCastOp = getDefiningCastOp(origRhs);
+    Type lhsElemType = maybeLhsCastOp ? getCastElemType(origLhs).value()
+                                      : getElemType(origLhs);
+    Type rhsElemType = maybeRhsCastOp ? getCastElemType(origRhs).value()
+                                      : getElemType(origRhs);
     Type outElemType = getElemType(origOut);
 
     if (!lhsElemType || !rhsElemType || !outElemType) {
@@ -310,13 +365,17 @@
     MatmulNarrowSizes narrowSizes =
         getMatmulNarrowSizes(origOut.getType().cast<ShapedType>());
     Location loc = matmulOp.getLoc();
-    TypeRange operandTypes = matmulOp->getOperandTypes();
-    Value encodedLhs = padAndSetEncoding(rewriter, loc, origLhs, user,
-                                         IREE::LinalgExt::EncodingRole::LHS,
-                                         operandTypes, narrowSizes);
-    Value encodedRhs = padAndSetEncoding(rewriter, loc, origRhs, user,
-                                         IREE::LinalgExt::EncodingRole::RHS,
-                                         operandTypes, narrowSizes);
+    SmallVector<Type> operandTypes(matmulOp->getOperandTypes());
+    operandTypes[0] =
+        cast<RankedTensorType>(operandTypes[0]).clone(lhsElemType);
+    operandTypes[1] =
+        cast<RankedTensorType>(operandTypes[1]).clone(rhsElemType);
+    Value encodedLhs = padAndSetEncoding(
+        rewriter, loc, origLhs, user, IREE::LinalgExt::EncodingRole::LHS,
+        operandTypes, narrowSizes, maybeLhsCastOp);
+    Value encodedRhs = padAndSetEncoding(
+        rewriter, loc, origRhs, user, IREE::LinalgExt::EncodingRole::RHS,
+        operandTypes, narrowSizes, maybeRhsCastOp);
     Value encodedOut = padAndSetEncoding(rewriter, loc, origOut, user,
                                          IREE::LinalgExt::EncodingRole::RESULT,
                                          operandTypes, narrowSizes);
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Utils.cpp b/compiler/src/iree/compiler/GlobalOptimization/Utils.cpp
new file mode 100644
index 0000000..a1a5859
--- /dev/null
+++ b/compiler/src/iree/compiler/GlobalOptimization/Utils.cpp
@@ -0,0 +1,57 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Utils.h"
+
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace GlobalOptimization {
+
+std::optional<CastOpInterface> getDefiningCastOp(Value input) {
+  auto castOp = input.getDefiningOp<CastOpInterface>();
+  if (castOp) {
+    return castOp;
+  }
+  auto genericOp = input.getDefiningOp<linalg::GenericOp>();
+  if (!genericOp || genericOp.getNumDpsInputs() != 1 ||
+      genericOp.getNumDpsInits() != 1 ||
+      genericOp.getBody()->getOperations().size() != 2 ||
+      !isElementwise(genericOp)) {
+    return std::nullopt;
+  }
+  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+  castOp = yieldOp->getOperand(0).getDefiningOp<CastOpInterface>();
+  if (!castOp) {
+    return std::nullopt;
+  }
+  Value castIn = castOp->getOperand(0);
+  if (castIn.isa<BlockArgument>() &&
+      castIn.cast<BlockArgument>().getArgNumber() != 0) {
+    return std::nullopt;
+  }
+  return castOp;
+}
+
+std::optional<Type> getCastElemType(Value input) {
+  std::optional<CastOpInterface> castOp = getDefiningCastOp(input);
+  if (!castOp) {
+    return std::nullopt;
+  }
+  Type castSrcElemType = getElementTypeOrSelf(castOp.value()->getOperand(0));
+  if (isa<arith::ExtUIOp>(castOp.value())) {
+    int64_t bitWidth = castSrcElemType.getIntOrFloatBitWidth();
+    return IntegerType::get(castOp->getContext(), bitWidth,
+                            IntegerType::SignednessSemantics::Unsigned);
+  }
+  return castSrcElemType;
+}
+
+} // namespace GlobalOptimization
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Utils.h b/compiler/src/iree/compiler/GlobalOptimization/Utils.h
new file mode 100644
index 0000000..03987bb
--- /dev/null
+++ b/compiler/src/iree/compiler/GlobalOptimization/Utils.h
@@ -0,0 +1,35 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#ifndef IREE_COMPILER_GLOBALOPTIMIZATION_UTILS_H_
+#define IREE_COMPILER_GLOBALOPTIMIZATION_UTILS_H_
+
+#include <optional>
+
+namespace mlir {
+class Type;
+class Value;
+class CastOpInterface;
+
+namespace iree_compiler {
+namespace GlobalOptimization {
+
+// If the producer is a CastOpInterface, or a linalg::GenericOp that performs
+// only a CastOpInterface on its input, return the CastOpInterface op.
+// Otherwise, return std::nullopt.
+//
+// **Note: If the CastOpInterface has been generalized, the return Operation
+//         is the body CastOpInterface op, not the linalg::GenericOp.
+std::optional<CastOpInterface> getDefiningCastOp(Value input);
+
+// Returns the source element type of the defining CastOpInterface of `input`,
+// if there is one. Otherwise return std::nullopt.
+std::optional<Type> getCastElemType(Value input);
+
+} // namespace GlobalOptimization
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_GLOBALOPTIMIZATION_UTILS_H_
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/set_encoding.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/set_encoding.mlir
index e675c55..5c7495c 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/test/set_encoding.mlir
+++ b/compiler/src/iree/compiler/GlobalOptimization/test/set_encoding.mlir
@@ -669,3 +669,89 @@
 // CHECK-NOT:     unset_encoding
 // CHECK:         linalg.matmul
 // CHECK:         linalg.batch_matmul
+
+// -----
+
+func.func @batch_matmul_casted_i8i8i32(%arg0 : tensor<64x100x250xi8>, %arg1 : tensor<64x250x500xi8>,
+      %arg2 : tensor<64x100x500xi32>) -> tensor<64x100x500xi32> {
+  %0 = tensor.empty() : tensor<64x250x500xi32>
+  %casted0 = arith.extui %arg0 : tensor<64x100x250xi8> to tensor<64x100x250xi32>
+  %casted1 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, 
+                                              affine_map<(d0, d1, d2) -> (d0, d1, d2)>], 
+                              iterator_types = ["parallel", "parallel", "parallel"]} 
+                              ins(%arg1 : tensor<64x250x500xi8>) 
+                              outs(%0 : tensor<64x250x500xi32>) {
+  ^bb0(%in: i8, %out: i32):
+      %2 = arith.extsi %in : i8 to i32
+      linalg.yield %2 : i32
+  } -> tensor<64x250x500xi32>
+  %1 = linalg.batch_matmul ins(%casted0, %casted1 : tensor<64x100x250xi32>, tensor<64x250x500xi32>)
+      outs(%arg2 : tensor<64x100x500xi32>) -> tensor<64x100x500xi32>
+  return %1 : tensor<64x100x500xi32>
+}
+//  CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (-s1 + (s1 ceildiv s0) * s0)>
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (-s1 + (s1 ceildiv s0) * s0 + 250)>
+//  CHECK-DAG: #[[MAP2:.+]] = affine_map<()[s0, s1] -> (-s1 + (s1 ceildiv s0) * s0 + 100)>
+//  CHECK-DAG: #[[MAP3:.+]] = affine_map<()[s0, s1] -> (-s1 + (s1 ceildiv s0) * s0 + 64)>
+//  CHECK-DAG: #[[MAP4:.+]] = affine_map<()[s0, s1] -> (-s1 + (s1 ceildiv s0) * s0 + 500)>
+//  CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+//      CHECK: func @batch_matmul_casted_i8i8i32(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<64x100x250xi8>
+// CHECK-SAME:     %[[ARG1:.+]]: tensor<64x250x500xi8>
+// CHECK-SAME:     %[[ARG2:.+]]: tensor<64x100x500xi32>
+//  CHECK-DAG:     %[[C64:.+]] = arith.constant 64 : index
+//  CHECK-DAG:     %[[C100:.+]] = arith.constant 100 : index
+//  CHECK-DAG:     %[[C250:.+]] = arith.constant 250 : index
+//  CHECK-DAG:     %[[C500:.+]] = arith.constant 500 : index
+//      CHECK:   %[[LHS_TILE_SIZE:.+]]:3 = iree_linalg_ext.upper_bound_tile_size tensor<64x100x250xi8, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = LHS, element_types = [ui8, i8, i32]>> -> index, index, index
+//      CHECK:   %[[LHS_PADDING_SIZE0:.+]] = affine.apply #[[MAP]]()[%[[LHS_TILE_SIZE]]#0, %[[C64]]]
+//      CHECK:   %[[LHS_PADDING_SIZE1:.+]] = affine.apply #[[MAP]]()[%[[LHS_TILE_SIZE]]#1, %[[C100]]]
+//      CHECK:   %[[LHS_PADDING_SIZE2:.+]] = affine.apply #[[MAP]]()[%[[LHS_TILE_SIZE]]#2, %[[C250]]]
+//      CHECK:   %[[LHS_PAD:.+]] = tensor.pad %[[ARG0]] low[0, 0, 0] high[%[[LHS_PADDING_SIZE0]], %[[LHS_PADDING_SIZE1]], %[[LHS_PADDING_SIZE2]]]
+//      CHECK:       tensor<64x100x250xi8> to tensor<?x?x?xi8>
+//  CHECK-DAG:   %[[LHS_DIM0:.+]] = affine.apply #[[MAP1]]()[%[[LHS_TILE_SIZE]]#2, %[[C250]]]
+//  CHECK-DAG:   %[[LHS_DIM1:.+]] = affine.apply #[[MAP2]]()[%[[LHS_TILE_SIZE]]#1, %[[C100]]]
+//  CHECK-DAG:   %[[LHS_DIM2:.+]] = affine.apply #[[MAP3]]()[%[[LHS_TILE_SIZE]]#0, %[[C64]]]
+//      CHECK:   %[[LHS:.+]] = iree_linalg_ext.set_encoding %[[LHS_PAD]]
+// CHECK-SAME:       tensor<?x?x?xi8, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = LHS, element_types = [ui8, i8, i32], original_type = tensor<64x100x250xi8>>>
+//      CHECK:   %[[INIT_LHS_CAST:.+]] = tensor.empty(%[[LHS_DIM2]], %[[LHS_DIM1]], %[[LHS_DIM0]]) : tensor<?x?x?xi32, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = LHS, element_types = [ui8, i8, i32], original_type = tensor<64x100x250xi8>>>
+//      CHECK:   %[[LHS_CASTED:.+]] = linalg.generic {indexing_maps = [#[[MAP5]], #[[MAP5]]], iterator_types = ["parallel", "parallel", "parallel"]}
+// CHECK-SAME:       ins(%[[LHS]] : tensor<?x?x?xi8, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = LHS, element_types = [ui8, i8, i32], original_type = tensor<64x100x250xi8>>>)
+// CHECK-SAME:       outs(%[[INIT_LHS_CAST]] : tensor<?x?x?xi32, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = LHS, element_types = [ui8, i8, i32], original_type = tensor<64x100x250xi8>>>)
+// CHECK-NEXT:   ^bb0(%[[LHS_ARG_IN:.+]]: i8, %[[LHS_ARG_OUT:.+]]: i32):
+// CHECK-NEXT:   %[[LHS_CAST_OP:.+]] = arith.extui %[[LHS_ARG_IN]] : i8 to i32
+// CHECK-NEXT:   linalg.yield %[[LHS_CAST_OP]] : i32
+// CHECK-NEXT:   -> tensor<?x?x?xi32, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = LHS, element_types = [ui8, i8, i32], original_type = tensor<64x100x250xi8>>>
+//      CHECK:   %[[RHS_TILE_SIZE:.+]]:3 = iree_linalg_ext.upper_bound_tile_size tensor<64x250x500xi8, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = RHS, element_types = [ui8, i8, i32]>> -> index, index, index
+//      CHECK:   %[[RHS_PADDING_SIZE0:.+]] = affine.apply #[[MAP]]()[%[[RHS_TILE_SIZE]]#0, %[[C64]]]
+//      CHECK:   %[[RHS_PADDING_SIZE1:.+]] = affine.apply #[[MAP]]()[%[[RHS_TILE_SIZE]]#1, %[[C250]]]
+//      CHECK:   %[[RHS_PADDING_SIZE2:.+]] = affine.apply #[[MAP]]()[%[[RHS_TILE_SIZE]]#2, %[[C500]]]
+//      CHECK:   %[[RHS_PAD:.+]] = tensor.pad %[[ARG1]] low[0, 0, 0] high[%[[RHS_PADDING_SIZE0]], %[[RHS_PADDING_SIZE1]], %[[RHS_PADDING_SIZE2]]]
+//      CHECK:       tensor<64x250x500xi8> to tensor<?x?x?xi8>
+//  CHECK-DAG:   %[[RHS_DIM0:.+]] = affine.apply #[[MAP4]]()[%[[RHS_TILE_SIZE]]#2, %[[C500]]]
+//  CHECK-DAG:   %[[RHS_DIM1:.+]] = affine.apply #[[MAP1]]()[%[[RHS_TILE_SIZE]]#1, %[[C250]]]
+//  CHECK-DAG:   %[[RHS_DIM2:.+]] = affine.apply #[[MAP3]]()[%[[RHS_TILE_SIZE]]#0, %[[C64]]]
+//      CHECK:   %[[RHS:.+]] = iree_linalg_ext.set_encoding %[[RHS_PAD]]
+// CHECK-SAME:       tensor<?x?x?xi8, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = RHS, element_types = [ui8, i8, i32], original_type = tensor<64x250x500xi8>>>
+//      CHECK:   %[[INIT_RHS_CAST:.+]] = tensor.empty(%[[RHS_DIM2]], %[[RHS_DIM1]], %[[RHS_DIM0]]) : tensor<?x?x?xi32, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = RHS, element_types = [ui8, i8, i32], original_type = tensor<64x250x500xi8>>>
+//      CHECK:   %[[RHS_CASTED:.+]] = linalg.generic {indexing_maps = [#[[MAP5]], #[[MAP5]]], iterator_types = ["parallel", "parallel", "parallel"]}
+// CHECK-SAME:       ins(%[[RHS]] : tensor<?x?x?xi8, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = RHS, element_types = [ui8, i8, i32], original_type = tensor<64x250x500xi8>>>)
+// CHECK-SAME:       outs(%[[INIT_RHS_CAST]] : tensor<?x?x?xi32, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = RHS, element_types = [ui8, i8, i32], original_type = tensor<64x250x500xi8>>>)
+// CHECK-NEXT:   ^bb0(%[[RHS_ARG_IN:.+]]: i8, %[[RHS_ARG_OUT:.+]]: i32):
+// CHECK-NEXT:   %[[RHS_CAST_OP:.+]] = arith.extsi %[[RHS_ARG_IN]] : i8 to i32
+// CHECK-NEXT:   linalg.yield %[[RHS_CAST_OP]] : i32
+// CHECK-NEXT:   -> tensor<?x?x?xi32, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = RHS, element_types = [ui8, i8, i32], original_type = tensor<64x250x500xi8>>>
+//      CHECK:   %[[OUTS_TILE_SIZE:.+]]:3 = iree_linalg_ext.upper_bound_tile_size tensor<64x100x500xi32, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = RESULT, element_types = [ui8, i8, i32]>> -> index, index, index
+//      CHECK:   %[[OUTS_PADDING_SIZE0:.+]] = affine.apply #[[MAP]]()[%[[OUTS_TILE_SIZE]]#0, %[[C64]]]
+//      CHECK:   %[[OUTS_PADDING_SIZE1:.+]] = affine.apply #[[MAP]]()[%[[OUTS_TILE_SIZE]]#1, %[[C100]]]
+//      CHECK:   %[[OUTS_PADDING_SIZE2:.+]] = affine.apply #[[MAP]]()[%[[OUTS_TILE_SIZE]]#2, %[[C500]]]
+//      CHECK:   %[[OUTS_PAD:.+]] = tensor.pad %[[ARG2]] low[0, 0, 0] high[%[[OUTS_PADDING_SIZE0]], %[[OUTS_PADDING_SIZE1]], %[[OUTS_PADDING_SIZE2]]]
+//      CHECK:       tensor<64x100x500xi32> to tensor<?x?x?xi32>
+//      CHECK:   %[[OUTS:.+]] = iree_linalg_ext.set_encoding %[[OUTS_PAD]]
+// CHECK-SAME:       tensor<?x?x?xi32, #iree_linalg_ext.encoding<user = BATCH_MATMUL, role = RESULT, element_types = [ui8, i8, i32], original_type = tensor<64x100x500xi32>>>
+//      CHECK:   %[[BATCH_MATMUL:.+]] = linalg.batch_matmul
+// CHECK-SAME:       ins(%[[LHS_CASTED]], %[[RHS_CASTED]] :
+// CHECK-SAME:       outs(%[[OUTS]] :
+//      CHECK:   %[[RESULT_PADDED:.+]] = iree_linalg_ext.unset_encoding %[[BATCH_MATMUL]]
+//      CHECK:   %[[RESULT:.+]] = tensor.extract_slice %[[RESULT_PADDED]][0, 0, 0] [64, 100, 500] [1, 1, 1]
+//      CHECK:   return %[[RESULT]]