diff --git a/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnBuffers.cpp b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnBuffers.cpp
index 50c240c..5c382dd 100644
--- a/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnBuffers.cpp
+++ b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnBuffers.cpp
@@ -637,7 +637,7 @@
     Value startIndex = rewriter.create<ConstantIndexOp>(
         loc, std::get<0>(it.value()).getZExtValue());
     offsets.push_back(startIndex);
-    Value size = rewriter.create<DimOp>(loc, resultBuffers[0], it.index());
+    Value size = rewriter.create<DimOp>(loc, inputBuffers[0], it.index());
     sizes.push_back(size);
     Value stride = rewriter.create<ConstantIndexOp>(
         loc, std::get<1>(it.value()).getZExtValue() + 1);
diff --git a/iree/compiler/Conversion/HLOToLinalg/test/pad.mlir b/iree/compiler/Conversion/HLOToLinalg/test/pad.mlir
index fe3309d..a42a079 100644
--- a/iree/compiler/Conversion/HLOToLinalg/test/pad.mlir
+++ b/iree/compiler/Conversion/HLOToLinalg/test/pad.mlir
@@ -1,12 +1,9 @@
-// RUN: iree-opt -split-input-file -iree-codegen-hlo-to-linalg-on-buffers %s | IreeFileCheck %s
+// RUN: iree-opt -split-input-file -iree-codegen-hlo-to-linalg-on-buffers -canonicalize %s | IreeFileCheck %s
 
 module {
-  // CHECK_LABEL: @pad_cst
   func @pad_cst() {
     %c0 = constant 0 : index
     %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<12x4xf32>
-    // CHECK: linalg.fill
-    // CHECK: linalg.copy
     %1 = constant dense<0.0> : tensor<f32>
     %2 = "mhlo.pad"(%0, %1) {
       edge_padding_high = dense<[2, 3]> : tensor<2xi64>,
@@ -21,17 +18,21 @@
     hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer", access="Write"
   }
 }
+// CHECK_LABEL: @pad_cst
+//   CHECK-DAG: %[[CST:.+]] = constant 0.000000e+00 : f32
+//   CHECK-DAG: %[[OUT:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<18x12xf32>
+//   CHECK-DAG: %[[IN:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<12x4xf32>
+//       CHECK: linalg.fill(%[[OUT]], %[[CST]])
+//       CHECK: %[[SUBVIEW:.+]] = subview %[[OUT]][4, 5] [12, 4] [1, 1]
+//       CHECK: linalg.copy(%[[IN]], %[[SUBVIEW]])
 
 // -----
 
 module {
-  // CHECK_LABEL: @pad_memref
   func @pad_memref() {
     %c0 = constant 0 : index
     %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<12x4xf32>
     %1 = hal.interface.load.tensor @legacy_io::@arg1, offset = %c0 : tensor<f32>
-    // CHECK: linalg.fill
-    // CHECK: linalg.copy
     %2 = "mhlo.pad"(%0, %1) {
       edge_padding_high = dense<[2, 3]> : tensor<2xi64>,
       edge_padding_low = dense<[4, 5]> : tensor<2xi64>,
@@ -46,16 +47,21 @@
     hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
   }
 }
+// CHECK_LABEL: @pad_memref
+//   CHECK-DAG: %[[OUT:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<18x12xf32>
+//   CHECK-DAG: %[[IN:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<12x4xf32>
+//   CHECK-DAG: %[[PAD_BUF:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg1} : memref<f32>
+//       CHECK: %[[PAD_VAL:.+]] = load %[[PAD_BUF]][] : memref<f32>
+//       CHECK: linalg.fill(%[[OUT]], %[[PAD_VAL]])
+//       CHECK: %[[SUBVIEW:.+]] = subview %[[OUT]][4, 5] [12, 4] [1, 1]
+//       CHECK: linalg.copy(%[[IN]], %[[SUBVIEW]])
 
 // -----
 
 module {
-  // CHECK_LABEL: @pad_no_op
   func @pad_no_op() {
     %c0 = constant 0 : index
     %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<12x4xf32>
-    // CHECK: linalg.fill
-    // CHECK: linalg.copy
     %1 = constant dense<0.0> : tensor<f32>
     %2 = "mhlo.pad"(%0, %1) {
       edge_padding_high = dense<0> : tensor<2xi64>,
@@ -70,17 +76,22 @@
     hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer", access="Write"
   }
 }
+// TODO(hanchung): Make it just a copy op.
+// CHECK_LABEL: @pad_no_op
+//   CHECK-DAG: %[[CST:.+]] = constant 0.000000e+00 : f32
+//   CHECK-DAG: %[[OUT:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<12x4xf32>
+//   CHECK-DAG: %[[IN:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<12x4xf32>
+//       CHECK: linalg.fill(%[[OUT]], %[[CST]])
+//       CHECK: %[[SUBVIEW:.+]] = subview %[[OUT]][0, 0] [12, 4] [1, 1]
+//       CHECK: linalg.copy(%[[IN]], %[[SUBVIEW]])
 
 // -----
 
 module {
-  // CHECK_LABEL: @cst_pad_cst
   func @cst_pad_cst() {
     %c0 = constant 0 : index
     %0 = constant dense<1.0> : tensor<12x4xf32>
     %1 = constant dense<0.0> : tensor<f32>
-    // CHECK: linalg.fill
-    // CHECK: linalg.fill
     %2 = "mhlo.pad"(%0, %1) {
       edge_padding_high = dense<[2, 3]> : tensor<2xi64>,
       edge_padding_low = dense<[4, 5]> : tensor<2xi64>,
@@ -93,3 +104,10 @@
     hal.interface.binding @ret0, set=0, binding=0, type="StorageBuffer", access="Write"
   }
 }
+// CHECK_LABEL: @cst_pad_cst
+//   CHECK-DAG: %[[ZERO:.+]] = constant 0.000000e+00 : f32
+//   CHECK-DAG: %[[ONE:.+]] = constant 1.000000e+00 : f32
+//   CHECK-DAG: %[[OUT:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<18x12xf32>
+//       CHECK: linalg.fill(%[[OUT]], %[[ZERO]])
+//       CHECK: %[[SUBVIEW:.+]] = subview %[[OUT]][4, 5] [12, 4] [1, 1]
+//       CHECK: linalg.fill(%[[SUBVIEW]], %[[ONE]])