diff --git a/iree/compiler/Conversion/CodegenUtils/FunctionUtils.h b/iree/compiler/Conversion/CodegenUtils/FunctionUtils.h
index adc2f53..6a4c81e 100644
--- a/iree/compiler/Conversion/CodegenUtils/FunctionUtils.h
+++ b/iree/compiler/Conversion/CodegenUtils/FunctionUtils.h
@@ -23,6 +23,16 @@
 /// Returns true if the given `func` is a kernel dispatch entry point.
 bool isEntryPoint(FuncOp func);
 
+/// Returns the attribute name used to record the binding associated with an
+/// iree.placeholder operation.
+inline const char* getBindingAttrName() { return "binding"; }
+
+/// Returns the attribute name used to record argument position in the (operand
+/// + result) list of shaped types of the dispatch region.
+inline const char* getOperandResultNumAttrName() {
+  return "operand_result_index";
+}
+
 }  // namespace iree_compiler
 }  // namespace mlir
 
diff --git a/iree/compiler/Conversion/HLOToLinalg/FusionOfTensorOps.cpp b/iree/compiler/Conversion/HLOToLinalg/FusionOfTensorOps.cpp
index 0808d42..ebf96aa 100644
--- a/iree/compiler/Conversion/HLOToLinalg/FusionOfTensorOps.cpp
+++ b/iree/compiler/Conversion/HLOToLinalg/FusionOfTensorOps.cpp
@@ -41,8 +41,8 @@
         reshapeOp.src().getDefiningOp<IREE::HAL::InterfaceLoadTensorOp>();
     if (!loadOp) return failure();
     rewriter.replaceOpWithNewOp<IREE::HAL::InterfaceLoadTensorOp>(
-        reshapeOp, reshapeOp.getResultType(), loadOp.binding(),
-        loadOp.offset());
+        reshapeOp, reshapeOp.getResultType(), loadOp.offset(),
+        loadOp.getAttrs());
     return success();
   }
 };
@@ -55,8 +55,9 @@
                                 PatternRewriter &rewriter) const override {
     auto reshapeOp = storeOp.operand().getDefiningOp<linalg::TensorReshapeOp>();
     if (!reshapeOp) return failure();
+    SmallVector<Value, 2> operands = {reshapeOp.src(), storeOp.offset()};
     rewriter.replaceOpWithNewOp<IREE::HAL::InterfaceStoreTensorOp>(
-        storeOp, reshapeOp.src(), storeOp.binding(), storeOp.offset());
+        storeOp, ArrayRef<Type>(), operands, storeOp.getAttrs());
     return success();
   }
 };
diff --git a/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnBuffers.cpp b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnBuffers.cpp
index 70a407f..87ed137 100644
--- a/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnBuffers.cpp
+++ b/iree/compiler/Conversion/HLOToLinalg/HLOToLinalgOnBuffers.cpp
@@ -22,6 +22,7 @@
 
 #include <cstddef>
 
+#include "iree/compiler/Conversion/CodegenUtils/FunctionUtils.h"
 #include "iree/compiler/Conversion/HLOToLinalg/Passes.h"
 #include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/IREE/IR/IREEDialect.h"
@@ -1348,7 +1349,10 @@
     // annotation is carried over if exists.
     auto phOp = rewriter.create<IREE::PlaceholderOp>(
         loadOp.getLoc(), bufferType, "interface buffer");
-    phOp.setAttr("binding", loadOp.binding());
+    phOp.setAttr(getBindingAttrName(), loadOp.binding());
+    StringRef attrName = getOperandResultNumAttrName();
+    if (auto operandResultNumAttr = loadOp.getAttr(attrName))
+      phOp.setAttr(attrName, operandResultNumAttr);
     Value buffer = phOp.getResult();
 
     // If the result of the load is already mapped to a buffer, a copy is
@@ -1463,7 +1467,10 @@
   // annotation is carried over if exists.
   auto phOp = builder.create<IREE::PlaceholderOp>(op.getLoc(), bufferType,
                                                   "interface buffer");
-  phOp.setAttr("binding", op.binding());
+  phOp.setAttr(getBindingAttrName(), op.binding());
+  StringRef attrName = getOperandResultNumAttrName();
+  if (Attribute operandResultNumAttr = op.getAttr(attrName))
+    phOp.setAttr(attrName, operandResultNumAttr);
   Value buffer = phOp;
   outputBufferMap[op] = buffer;
 
@@ -1576,7 +1583,8 @@
   target.addIllegalOp<linalg::TensorReshapeOp>();
   target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
       Optional<ConversionTarget::DynamicLegalityCallbackFn>([](Operation *op) {
-        // The generated structured Linalg ops should have buffer semantics.
+        // The generated structured Linalg ops should have buffer
+        // semantics.
         if (auto linalgOp = dyn_cast<linalg::LinalgOp>(op)) {
           return linalgOp.hasBufferSemantics();
         }
diff --git a/iree/compiler/Conversion/HLOToLinalg/test/fusion.mlir b/iree/compiler/Conversion/HLOToLinalg/test/fusion.mlir
index a82179b..bb090b8 100644
--- a/iree/compiler/Conversion/HLOToLinalg/test/fusion.mlir
+++ b/iree/compiler/Conversion/HLOToLinalg/test/fusion.mlir
@@ -38,7 +38,7 @@
     return
   }
   hal.interface @legacy_io attributes {sym_visiblity = "private"} {
-    hal.interface.binding @ret0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=0, type="StorageBuffer", access="Write|Discard"
   }
 }
 
diff --git a/iree/compiler/Conversion/HLOToLinalg/test/linalg_tensor_to_buffer.mlir b/iree/compiler/Conversion/HLOToLinalg/test/linalg_tensor_to_buffer.mlir
index 2bab77a..9363815 100644
--- a/iree/compiler/Conversion/HLOToLinalg/test/linalg_tensor_to_buffer.mlir
+++ b/iree/compiler/Conversion/HLOToLinalg/test/linalg_tensor_to_buffer.mlir
@@ -5,29 +5,35 @@
 module {
   func @element_wise() {
     %c0 = constant 0 : index
-    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<2x2xf32>
-    %1 = hal.interface.load.tensor @legacy_io::@arg1, offset = %c0 : tensor<2x2xf32>
-    %2 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]} %0, %1 {
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<2x2xf32>
+    %1 = hal.interface.load.tensor @legacy_io::@arg1, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<2x2xf32>
+    %2 = linalg.generic
+      {args_in = 2 : i64, args_out = 1 : i64,
+       indexing_maps = [#map0, #map0, #map0],
+       iterator_types = ["parallel", "parallel"]} %0, %1 {
     ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
       %3 = addf %arg3, %arg4 : f32
       linalg.yield %3 : f32
     }: tensor<2x2xf32>, tensor<2x2xf32> -> tensor<2x2xf32>
-    hal.interface.store.tensor %2, @legacy_io::@ret0, offset = %c0 : tensor<2x2xf32>
+    hal.interface.store.tensor %2, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 2 : i32} : tensor<2x2xf32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
-    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
-    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
-    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer",
+      access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer",
+      access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer",
+      access="Write"
   }
 }
 // CHECK-LABEL: func @element_wise
-//   CHECK-DAG: %[[ARG2:.*]] = iree.placeholder
-//  CHECK-SAME:   {binding = @legacy_io::@ret0}
-//   CHECK-DAG: %[[ARG0:.*]] = iree.placeholder
-//  CHECK-SAME:   {binding = @legacy_io::@arg0}
-//   CHECK-DAG: %[[ARG1:.*]] = iree.placeholder
-//  CHECK-SAME:   {binding = @legacy_io::@arg1}
+//   CHECK-DAG: %[[ARG2:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 2 : i32}
+//   CHECK-DAG: %[[ARG0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
+//   CHECK-DAG: %[[ARG1:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg1, operand_result_index = 1 : i32}
 //   CHECK-NOT: hal.interface.load.tensor
 //       CHECK: linalg.generic
 //  CHECK-SAME:   %[[ARG0]], %[[ARG1]], %[[ARG2]]
@@ -45,8 +51,12 @@
 module {
   func @indexed_generic() {
     %c0 = constant 0 : index
-    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<2x2xi32>
-    %1 = linalg.indexed_generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]} %0 {
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<2x2xi32>
+    %1 = linalg.indexed_generic
+      {args_in = 1 : i64, args_out = 1 : i64,
+       indexing_maps = [#map0, #map0],
+       iterator_types = ["parallel", "parallel"]} %0 {
     ^bb0(%arg2: index, %arg3: index, %arg4: i32):       // no predecessors
       %2 = index_cast %arg2 : index to i32
       %3 = index_cast %arg3 : index to i32
@@ -54,19 +64,20 @@
       %5 = addi %4, %3 : i32
       linalg.yield %5 : i32
     }: tensor<2x2xi32> -> tensor<2x2xi32>
-    hal.interface.store.tensor %1, @legacy_io::@ret0, offset = %c0 : tensor<2x2xi32>
+    hal.interface.store.tensor %1, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<2x2xi32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
-    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
-    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer", access="Write"
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer",
+      access="Read"
+    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer",
+      access="Write"
   }
 }
 //      CHECK: func @indexed_generic
-//  CHECK-DAG: %[[RET0:.*]] = iree.placeholder
-// CHECK-SAME:   {binding = @legacy_io::@ret0}
-//  CHECK-DAG: %[[ARG0:.*]] = iree.placeholder
-// CHECK-SAME:   {binding = @legacy_io::@arg0}
+//  CHECK-DAG: %[[RET0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 1 : i32}
+//  CHECK-DAG: %[[ARG0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
 //  CHECK-NOT: hal.interface.load.tensor
 //      CHECK: linalg.indexed_generic
 // CHECK-SAME:   %[[ARG0]], %[[RET0]]
@@ -90,23 +101,23 @@
 module {
   func @reshape_arg_result() {
     %c0 = constant 0 : index
-    %0 = hal.interface.load.tensor @legacy_io::@arg0,
-                                   offset = %c0 : tensor<5xf32>
-    %1 = hal.interface.load.tensor @legacy_io::@arg1,
-                                   offset = %c0 : tensor<5xf32>
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<5xf32>
+    %1 = hal.interface.load.tensor @legacy_io::@arg1, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<5xf32>
     %2 = linalg.tensor_reshape %0 [#map0] : tensor<5xf32> into tensor<5x1xf32>
     %3 = linalg.tensor_reshape %1 [#map0] : tensor<5xf32> into tensor<1x5xf32>
     %4 = linalg.generic
-           {args_in = 2 : i64, args_out = 1 : i64,
-            indexing_maps = [#map1, #map2, #map0],
-            iterator_types = ["parallel", "parallel"]} %2, %3 {
+      {args_in = 2 : i64, args_out = 1 : i64,
+       indexing_maps = [#map1, #map2, #map0],
+       iterator_types = ["parallel", "parallel"]} %2, %3 {
          ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
            %5 = addf %arg3, %arg4 : f32
            linalg.yield %5 : f32
          }: tensor<5x1xf32>, tensor<1x5xf32> -> tensor<5x5xf32>
     %6 = linalg.tensor_reshape %4 [#map0] : tensor<5x5xf32> into tensor<25xf32>
-    hal.interface.store.tensor %6, @legacy_io::@ret0,
-                                   offset = %c0 : tensor<25xf32>
+    hal.interface.store.tensor %6, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 2 : i32} : tensor<25xf32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
@@ -122,13 +133,10 @@
 //   CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d0, 0)>
 //   CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1) -> (0, d1)>
 //       CHECK: func @reshape_arg_result
-//   CHECK-DAG:   %[[RET0:.*]] = iree.placeholder
-//  CHECK-SAME:     binding = @legacy_io::@ret0
+//   CHECK-DAG:   %[[RET0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 2 : i32}
 //   CHECK-DAG:   %[[RESULT:.*]] = linalg.reshape %[[RET0]] [#[[MAP0]]]
-//   CHECK-DAG:   %[[ARG0:.*]] = iree.placeholder
-//  CHECK-SAME:     binding = @legacy_io::@arg0
-//   CHECK-DAG:   %[[ARG1:.*]] = iree.placeholder
-//  CHECK-SAME:     binding = @legacy_io::@arg1
+//   CHECK-DAG:   %[[ARG0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
+//   CHECK-DAG:   %[[ARG1:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg1, operand_result_index = 1 : i32}
 //   CHECK-DAG:   %[[LHS:.*]] = linalg.reshape %[[ARG0]] [#[[MAP0]]]
 //   CHECK-DAG:   %[[RHS:.*]] = linalg.reshape %[[ARG1]] [#[[MAP0]]]
 //       CHECK:   linalg.generic
@@ -141,11 +149,11 @@
 module {
   func @reshape_only() {
     %c0 = constant 0 : index
-    %0 = hal.interface.load.tensor @legacy_io::@arg0,
-                                   offset = %c0 : tensor<5x5xf32>
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<5x5xf32>
     %1 = linalg.tensor_reshape %0 [#map0] : tensor<5x5xf32> into tensor<25xf32>
-    hal.interface.store.tensor %1, @legacy_io::@ret0,
-                                   offset = %c0 : tensor<25xf32>
+    hal.interface.store.tensor %1, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<25xf32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
@@ -157,11 +165,9 @@
 }
 //       CHECK: #[[MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 //       CHECK: func @reshape_only
-//   CHECK-DAG:   %[[RET0:.*]] = iree.placeholder
-//  CHECK-SAME:     binding = @legacy_io::@ret0
+//   CHECK-DAG:   %[[RET0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 1 : i32}
 //   CHECK-DAG:   %[[RESULT:.*]] = linalg.reshape %[[RET0]] [#[[MAP0]]]
-//   CHECK-DAG:   %[[ARG0:.*]] = iree.placeholder
-//  CHECK-SAME:     binding = @legacy_io::@arg0
+//   CHECK-DAG:   %[[ARG0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
 //       CHECK:   linalg.copy(%[[ARG0]], %[[RESULT]])
 
 // -----
@@ -171,34 +177,36 @@
 module {
   func @store_value_twice() {
     %c0 = constant 0 : index
-    %0 = hal.interface.load.tensor @legacy_io::@arg0,
-                                   offset = %c0 : tensor<2x4xf32>
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<2x4xf32>
     %1 = linalg.generic
-           {args_in = 1 : i64, args_out = 1 : i64,
-            indexing_maps = [#map0, #map0],
-            iterator_types = ["parallel", "parallel"]} %0 {
+      {args_in = 1 : i64, args_out = 1 : i64,
+       indexing_maps = [#map0, #map0],
+       iterator_types = ["parallel", "parallel"]} %0 {
     ^bb0(%arg0: f32):  // no predecessors
       %2 = tanh %arg0 : f32
       linalg.yield %2 : f32
     }: tensor<2x4xf32> -> tensor<2x4xf32>
-    hal.interface.store.tensor %1, @legacy_io::@ret0, offset = %c0 : tensor<2x4xf32>
-    hal.interface.store.tensor %1, @legacy_io::@ret1, offset = %c0 : tensor<2x4xf32>
+    hal.interface.store.tensor %1, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<2x4xf32>
+    hal.interface.store.tensor %1, @legacy_io::@ret1, offset = %c0
+      {operand_result_index = 2 : i32} : tensor<2x4xf32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
-    hal.interface.binding @arg0, set=0, binding=0,
-                                 type="StorageBuffer", access="Read"
-    hal.interface.binding @ret0, set=0, binding=1,
-                                 type="StorageBuffer", access="Write|Discard"
-    hal.interface.binding @ret1, set=0, binding=2,
-                                 type="StorageBuffer", access="Write|Discard"
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer",
+      access="Read"
+    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer",
+      access="Write|Discard"
+    hal.interface.binding @ret1, set=0, binding=2, type="StorageBuffer",
+      access="Write|Discard"
   }
 }
 
 // CHECK-LABEL: func @store_value_twice
-//       CHECK:   %[[T0:.*]] = iree.placeholder {{.*}} @legacy_io::@ret0
-//       CHECK:   %[[T1:.*]] = iree.placeholder {{.*}} @legacy_io::@ret1
-//       CHECK:   %[[T2:.*]] = iree.placeholder {{.*}} @legacy_io::@arg0
+//   CHECK-DAG:   %[[T0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 1 : i32}
+//       CHECK:   %[[T1:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret1, operand_result_index = 2 : i32}
+//       CHECK:   %[[T2:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
 //       CHECK:   linalg.generic {{.*}} %[[T2]], %[[T0]]
 //       CHECK:   linalg.copy(%[[T0]], %[[T1]])
 
@@ -211,36 +219,39 @@
 module {
   func @store_reshape_src_and_result_0() {
     %c0 = constant 0 : index
-    %0 = hal.interface.load.tensor @legacy_io::@arg0,
-                                   offset = %c0 : tensor<2x4xf32>
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<2x4xf32>
     %1 = linalg.generic
-           {args_in = 1 : i64, args_out = 1 : i64,
-            indexing_maps = [#map0, #map0],
-            iterator_types = ["parallel", "parallel"]} %0 {
+      {args_in = 1 : i64, args_out = 1 : i64,
+       indexing_maps = [#map0, #map0],
+       iterator_types = ["parallel", "parallel"]} %0 {
     ^bb0(%arg0: f32):  // no predecessors
       %2 = tanh %arg0 : f32
       linalg.yield %2 : f32
     }: tensor<2x4xf32> -> tensor<2x4xf32>
-    %3 = linalg.tensor_reshape %1 [#map1, #map2] : tensor<2x4xf32> into tensor<1x2x4xf32>
-    hal.interface.store.tensor %3, @legacy_io::@ret1, offset = %c0 : tensor<1x2x4xf32>
-    hal.interface.store.tensor %1, @legacy_io::@ret0, offset = %c0 : tensor<2x4xf32>
+    %3 = linalg.tensor_reshape %1 [#map1, #map2]
+      : tensor<2x4xf32> into tensor<1x2x4xf32>
+    hal.interface.store.tensor %3, @legacy_io::@ret1, offset = %c0
+      {operand_result_index = 2 : i32} : tensor<1x2x4xf32>
+    hal.interface.store.tensor %1, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<2x4xf32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
-    hal.interface.binding @arg0, set=0, binding=0,
-                                 type="StorageBuffer", access="Read"
-    hal.interface.binding @ret0, set=0, binding=1,
-                                 type="StorageBuffer", access="Write|Discard"
-    hal.interface.binding @ret1, set=0, binding=2,
-                                 type="StorageBuffer", access="Write|Discard"
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer",
+      access="Read"
+    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer",
+      access="Write|Discard"
+    hal.interface.binding @ret1, set=0, binding=2, type="StorageBuffer",
+      access="Write|Discard"
   }
 }
 
 // CHECK-LABEL: func @store_reshape_src_and_result_0
-//       CHECK:   %[[T0:.*]] = iree.placeholder {{.*}} @legacy_io::@ret1
+//       CHECK:   %[[T0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret1, operand_result_index = 2 : i32}
 //       CHECK:   %[[T1:.*]] = linalg.reshape %[[T0]]
-//       CHECK:   %[[T2:.*]] = iree.placeholder {{.*}} @legacy_io::@ret0
-//       CHECK:   %[[T3:.*]] = iree.placeholder {{.*}} @legacy_io::@arg0
+//       CHECK:   %[[T2:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 1 : i32}
+//       CHECK:   %[[T3:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
 //       CHECK:   linalg.generic {{.*}} %[[T3]], %[[T1]]
 //       CHECK:   linalg.copy(%[[T1]], %[[T2]])
 //       CHECK:   return
@@ -254,35 +265,38 @@
 module {
   func @store_reshape_src_and_result_1() {
     %c0 = constant 0 : index
-    %0 = hal.interface.load.tensor @legacy_io::@arg0,
-                                   offset = %c0 : tensor<2x4xf32>
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<2x4xf32>
     %1 = linalg.generic
-           {args_in = 1 : i64, args_out = 1 : i64,
-            indexing_maps = [#map0, #map0],
-            iterator_types = ["parallel", "parallel"]} %0 {
+      {args_in = 1 : i64, args_out = 1 : i64,
+       indexing_maps = [#map0, #map0],
+       iterator_types = ["parallel", "parallel"]} %0 {
     ^bb0(%arg0: f32):  // no predecessors
       %2 = tanh %arg0 : f32
       linalg.yield %2 : f32
     }: tensor<2x4xf32> -> tensor<2x4xf32>
-    %3 = linalg.tensor_reshape %1 [#map1, #map2] : tensor<2x4xf32> into tensor<1x2x4xf32>
-    hal.interface.store.tensor %1, @legacy_io::@ret0, offset = %c0 : tensor<2x4xf32>
-    hal.interface.store.tensor %3, @legacy_io::@ret1, offset = %c0 : tensor<1x2x4xf32>
+    %3 = linalg.tensor_reshape %1 [#map1, #map2]
+      : tensor<2x4xf32> into tensor<1x2x4xf32>
+    hal.interface.store.tensor %1, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<2x4xf32>
+    hal.interface.store.tensor %3, @legacy_io::@ret1, offset = %c0
+      {operand_result_index = 2 : i32} : tensor<1x2x4xf32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
-    hal.interface.binding @arg0, set=0, binding=0,
-                                 type="StorageBuffer", access="Read"
-    hal.interface.binding @ret0, set=0, binding=1,
-                                 type="StorageBuffer", access="Write|Discard"
-    hal.interface.binding @ret1, set=0, binding=2,
-                                 type="StorageBuffer", access="Write|Discard"
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer",
+      access="Read"
+    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer",
+      access="Write|Discard"
+    hal.interface.binding @ret1, set=0, binding=2, type="StorageBuffer",
+    access="Write|Discard"
   }
 }
 
 // CHECK-LABEL: func @store_reshape_src_and_result_1
-//       CHECK:   %[[T0:.*]] = iree.placeholder {{.*}} @legacy_io::@ret0
-//       CHECK:   %[[T1:.*]] = iree.placeholder {{.*}} @legacy_io::@ret1
-//       CHECK:   %[[T2:.*]] = iree.placeholder {{.*}} @legacy_io::@arg0
+//   CHECK-DAG:   %[[T0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 1 : i32}
+//   CHECK-DAG:   %[[T1:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret1, operand_result_index = 2 : i32}
+//   CHECK-DAG:   %[[T2:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
 //       CHECK:   linalg.generic {{.*}} %[[T2]], %[[T0]]
 //       CHECK:   %[[T3:.*]] = linalg.reshape %[[T0]]
 //       CHECK:   linalg.copy(%[[T3]], %[[T1]])
@@ -297,42 +311,48 @@
 module {
   func @store_reshape_src_and_result_2() {
     %c0 = constant 0 : index
-    %0 = hal.interface.load.tensor @legacy_io::@arg0,
-                                   offset = %c0 : tensor<2x4xf32>
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<2x4xf32>
     %1 = linalg.generic
-           {args_in = 1 : i64, args_out = 1 : i64,
-            indexing_maps = [#map0, #map0],
-            iterator_types = ["parallel", "parallel"]} %0 {
+      {args_in = 1 : i64, args_out = 1 : i64,
+       indexing_maps = [#map0, #map0],
+       iterator_types = ["parallel", "parallel"]} %0 {
     ^bb0(%arg0: f32):  // no predecessors
       %2 = tanh %arg0 : f32
       linalg.yield %2 : f32
     }: tensor<2x4xf32> -> tensor<2x4xf32>
-    %3 = linalg.tensor_reshape %1 [#map1, #map2] : tensor<2x4xf32> into tensor<1x2x4xf32>
-    %4 = linalg.tensor_reshape %1 [#map1, #map2] : tensor<2x4xf32> into tensor<1x2x4xf32>
-    %5 = linalg.tensor_reshape %1 [#map1, #map2] : tensor<2x4xf32> into tensor<1x2x4xf32>
-    hal.interface.store.tensor %3, @legacy_io::@ret0, offset = %c0 : tensor<1x2x4xf32>
-    hal.interface.store.tensor %4, @legacy_io::@ret1, offset = %c0 : tensor<1x2x4xf32>
-    hal.interface.store.tensor %5, @legacy_io::@ret2, offset = %c0 : tensor<1x2x4xf32>
+    %3 = linalg.tensor_reshape %1 [#map1, #map2]
+      : tensor<2x4xf32> into tensor<1x2x4xf32>
+    %4 = linalg.tensor_reshape %1 [#map1, #map2]
+      : tensor<2x4xf32> into tensor<1x2x4xf32>
+    %5 = linalg.tensor_reshape %1 [#map1, #map2]
+      : tensor<2x4xf32> into tensor<1x2x4xf32>
+    hal.interface.store.tensor %3, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<1x2x4xf32>
+    hal.interface.store.tensor %4, @legacy_io::@ret1, offset = %c0
+      {operand_result_index = 2 : i32} : tensor<1x2x4xf32>
+    hal.interface.store.tensor %5, @legacy_io::@ret2, offset = %c0
+      {operand_result_index = 3 : i32} : tensor<1x2x4xf32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
-    hal.interface.binding @arg0, set=0, binding=0,
-                                 type="StorageBuffer", access="Read"
-    hal.interface.binding @ret0, set=0, binding=1,
-                                 type="StorageBuffer", access="Write|Discard"
-    hal.interface.binding @ret1, set=0, binding=2,
-                                 type="StorageBuffer", access="Write|Discard"
-    hal.interface.binding @ret2, set=0, binding=3,
-                                 type="StorageBuffer", access="Write|Discard"
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer",
+      access="Read"
+    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer",
+      access="Write|Discard"
+    hal.interface.binding @ret1, set=0, binding=2, type="StorageBuffer",
+      access="Write|Discard"
+    hal.interface.binding @ret2, set=0, binding=3, type="StorageBuffer",
+      access="Write|Discard"
   }
 }
 
 // CHECK-LABEL: func @store_reshape_src_and_result_2
-//       CHECK:   %[[T0:.*]] = iree.placeholder {{.*}} @legacy_io::@ret0
+//   CHECK-DAG:   %[[T0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 1 : i32}
 //       CHECK:   %[[T1:.*]] = linalg.reshape %[[T0]]
-//       CHECK:   %[[T2:.*]] = iree.placeholder {{.*}} @legacy_io::@ret1
-//       CHECK:   %[[T3:.*]] = iree.placeholder {{.*}} @legacy_io::@ret2
-//       CHECK:   %[[T4:.*]] = iree.placeholder {{.*}} @legacy_io::@arg0
+//   CHECK-DAG:   %[[T2:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret1, operand_result_index = 2 : i32}
+//   CHECK-DAG:   %[[T3:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret2, operand_result_index = 3 : i32}
+//   CHECK-DAG:   %[[T4:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
 //       CHECK:   linalg.generic {{.*}} %[[T4]], %[[T1]]
 //       CHECK:   linalg.copy(%[[T0]], %[[T2]])
 //       CHECK:   linalg.copy(%[[T0]], %[[T3]])
@@ -347,27 +367,37 @@
 module {
   func @edge_detect_sobel_operator_ex_dispatch_3() {
     %c0 = constant 0 : index
-    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<1x128x128x1xf32>
-    %1 = linalg.tensor_reshape %0 [#map0, #map1] : tensor<1x128x128x1xf32> into tensor<128x128xf32>
-    %2 = linalg.tensor_reshape %0 [#map0, #map1] : tensor<1x128x128x1xf32> into tensor<128x128xf32>
-    %3 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map2, #map2, #map2], iterator_types = ["parallel", "parallel"]} %1, %2 {
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<1x128x128x1xf32>
+    %1 = linalg.tensor_reshape %0 [#map0, #map1]
+      : tensor<1x128x128x1xf32> into tensor<128x128xf32>
+    %2 = linalg.tensor_reshape %0 [#map0, #map1]
+      : tensor<1x128x128x1xf32> into tensor<128x128xf32>
+    %3 = linalg.generic
+      {args_in = 2 : i64, args_out = 1 : i64,
+       indexing_maps = [#map2, #map2, #map2],
+       iterator_types = ["parallel", "parallel"]} %1, %2 {
     ^bb0(%arg0: f32, %arg1: f32):  // no predecessors
       %5 = mulf %arg0, %arg1 : f32
       linalg.yield %5 : f32
     }: tensor<128x128xf32>, tensor<128x128xf32> -> tensor<128x128xf32>
-    %4 = linalg.tensor_reshape %3 [#map0, #map1] : tensor<128x128xf32> into tensor<1x128x128x1xf32>
-    hal.interface.store.tensor %4, @legacy_io::@ret0, offset = %c0 : tensor<1x128x128x1xf32>
+    %4 = linalg.tensor_reshape %3 [#map0, #map1]
+      : tensor<128x128xf32> into tensor<1x128x128x1xf32>
+    hal.interface.store.tensor %4, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<1x128x128x1xf32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
-    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
-    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer", access="Write|Discard"
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer",
+      access="Read"
+    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer",
+      access="Write|Discard"
   }
 }
 // CHECK-LABEL: func @edge_detect_sobel_operator_ex_dispatch_3
-//       CHECK:   %[[T0:.*]] = iree.placeholder {{.*}} @legacy_io::@ret0
+//       CHECK:   %[[T0:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0, operand_result_index = 1 : i32}
 //       CHECK:   %[[T1:.*]] = linalg.reshape %[[T0]]
-//       CHECK:   %[[T2:.*]] = iree.placeholder {{.*}} @legacy_io::@arg0
+//       CHECK:   %[[T2:.*]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
 //       CHECK:   %[[T3:.*]] = linalg.reshape %[[T2]]
 //       CHECK:   %[[T4:.*]] = linalg.reshape %[[T2]]
 //       CHECK:   linalg.generic {{.*}} %[[T3]], %[[T4]], %[[T1]]
@@ -383,35 +413,44 @@
   func @generic_reshape_reshape() {
     %c0 = constant 0 : index
     %cst = constant 0.000000e+00 : f32
-    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<1x1x1x1000xf32>
-    %1 = linalg.tensor_reshape %0 [#map0] : tensor<1x1x1x1000xf32> into tensor<1000xf32>
-    %2 = linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map1, #map1], iterator_types = ["parallel"]} %1 {
+    %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0
+      {operand_result_index = 0 : i32} : tensor<1x1x1x1000xf32>
+    %1 = linalg.tensor_reshape %0 [#map0]
+      : tensor<1x1x1x1000xf32> into tensor<1000xf32>
+    %2 = linalg.generic
+      {args_in = 1 : i64, args_out = 1 : i64,
+       indexing_maps = [#map1, #map1], iterator_types = ["parallel"]} %1 {
     ^bb0(%arg0: f32):  // no predecessors
       %5 = addf %arg0, %cst : f32
       linalg.yield %5 : f32
     }: tensor<1000xf32> -> tensor<1000xf32>
-    %3 = linalg.tensor_reshape %2 [#map0] : tensor<1000xf32> into tensor<1x1x1x1000xf32>
-    %4 = linalg.tensor_reshape %3 [#map2, #map3] : tensor<1x1x1x1000xf32> into tensor<1x1000xf32>
-    hal.interface.store.tensor %3, @legacy_io::@ret0, offset = %c0 : tensor<1x1x1x1000xf32>
-    hal.interface.store.tensor %4, @legacy_io::@ret1, offset = %c0 : tensor<1x1000xf32>
+    %3 = linalg.tensor_reshape %2 [#map0]
+      : tensor<1000xf32> into tensor<1x1x1x1000xf32>
+    %4 = linalg.tensor_reshape %3 [#map2, #map3]
+      : tensor<1x1x1x1000xf32> into tensor<1x1000xf32>
+    hal.interface.store.tensor %3, @legacy_io::@ret0, offset = %c0
+      {operand_result_index = 1 : i32} : tensor<1x1x1x1000xf32>
+    hal.interface.store.tensor %4, @legacy_io::@ret1, offset = %c0
+      {operand_result_index = 2 : i32} : tensor<1x1000xf32>
     return
   }
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
-    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
-    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer", access="Write|Discard"
-    hal.interface.binding @ret1, set=0, binding=2, type="StorageBuffer", access="Write|Discard"
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer",
+      access="Read"
+    hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer",
+      access="Write|Discard"
+    hal.interface.binding @ret1, set=0, binding=2, type="StorageBuffer",
+      access="Write|Discard"
   }
 }
 
 // CHECK-LABEL: func @generic_reshape_reshape
 //       CHECK:   %[[RET0:.+]] = iree.placeholder
-//  CHECK-SAME:     @legacy_io::@ret0
+//  CHECK-SAME:     binding = @legacy_io::@ret0, operand_result_index = 1 : i32
 //       CHECK:   %[[RET0_RESHAPE:.+]] = linalg.reshape %[[RET0]]
 //  CHECK-SAME:     memref<1x1x1x1000xf32> into memref<1000xf32>
-//       CHECK:   %[[RET1:.+]] = iree.placeholder
-//  CHECK-SAME:     @legacy_io::@ret1
-//       CHECK:   %[[ARG0:.+]] = iree.placeholder
-//  CHECK-SAME:     @legacy_io::@arg0
+//       CHECK:   %[[RET1:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret1, operand_result_index = 2 : i32}
+//       CHECK:   %[[ARG0:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0, operand_result_index = 0 : i32}
 //       CHECK:   %[[ARG0_RESHAPE:.+]] = linalg.reshape %[[ARG0]]
 //  CHECK-SAME:     memref<1x1x1x1000xf32> into memref<1000xf32>
 //       CHECK:   linalg.generic
diff --git a/iree/compiler/Conversion/HLOToLinalg/test/pipeline_test.mlir b/iree/compiler/Conversion/HLOToLinalg/test/pipeline_test.mlir
index c5b558b..90c4749 100644
--- a/iree/compiler/Conversion/HLOToLinalg/test/pipeline_test.mlir
+++ b/iree/compiler/Conversion/HLOToLinalg/test/pipeline_test.mlir
@@ -46,14 +46,10 @@
   }
 }
 // CHECK-LABEL: func @bug_2882_repro2
-//       CHECK:   %[[RET0:.+]] = iree.placeholder
-//  CHECK-SAME:     binding = @legacy_io::@ret0
-//       CHECK:   %[[RET1:.+]] = iree.placeholder
-//  CHECK-SAME:     binding = @legacy_io::@ret1
-//       CHECK:   %[[ARG1:.+]] = iree.placeholder
-//  CHECK-SAME:     binding = @legacy_io::@arg1
-//       CHECK:   %[[ARG0:.+]] = iree.placeholder
-//  CHECK-SAME:     binding = @legacy_io::@arg0
+//   CHECK-DAG:   %[[RET0:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0}
+//   CHECK-DAG:   %[[RET1:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret1}
+//   CHECK-DAG:   %[[ARG1:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg1}
+//   CHECK-DAG:   %[[ARG0:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0}
 //       CHECK:   linalg.generic
 //  CHECK-SAME:     %[[ARG0]], %[[ARG1]], %[[RET0]]
 //       CHECK:   linalg.copy(%[[RET0]], %[[RET1]])
diff --git a/iree/compiler/Conversion/HLOToLinalg/test/slice.mlir b/iree/compiler/Conversion/HLOToLinalg/test/slice.mlir
index 300c0ff..0c17311 100644
--- a/iree/compiler/Conversion/HLOToLinalg/test/slice.mlir
+++ b/iree/compiler/Conversion/HLOToLinalg/test/slice.mlir
@@ -4,7 +4,8 @@
   // CHECK_LABEL: @slice_whole_buffer
   //  CHECK-NOT: subview
   //      CHECK: linalg.copy
-  func @slice_whole_buffer() {
+  func @slice_whole_buffer()
+    attributes {signature = (tensor<3x4xi32>) -> (tensor<3x4xi32>)} {
     %c0 = constant 0 : index
     %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<3x4xi32>
     %1 = "mhlo.slice"(%0) {
@@ -38,7 +39,8 @@
   // CHECK-SAME:   [%[[ONE]], %[[ONE]]]
   // CHECK-SAME: : memref<3x4xi32> to memref<?x?xi32, #[[MAP]]>
   //      CHECK: linalg.copy
-  func @slice_whole_stride() {
+  func @slice_whole_stride()
+    attributes {signature = (tensor<3x4xi32>) -> (tensor<1x4xi32>)} {
     %c0 = constant 0 : index
     %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<3x4xi32>
     %1 = "mhlo.slice"(%0) {
@@ -72,7 +74,8 @@
   // CHECK-SAME:   [%[[ONE]], %[[ONE]]]
   // CHECK-SAME: : memref<3x4xi32> to memref<?x?xi32, #map0>
   //       CHECK: linalg.copy
-  func @slice_stride_part() {
+  func @slice_stride_part()
+    attributes {signature = (tensor<3x4xi32>) -> (tensor<1x2xi32>)} {
     %c0 = constant 0 : index
     %0 = hal.interface.load.tensor @legacy_io::@arg0, offset = %c0 : tensor<3x4xi32>
     %1 = "mhlo.slice"(%0) {
diff --git a/iree/compiler/Conversion/LinalgToLLVM/test/matmul_vectorization.mlir b/iree/compiler/Conversion/LinalgToLLVM/test/matmul_vectorization.mlir
index 7d47f47..123a1e0 100644
--- a/iree/compiler/Conversion/LinalgToLLVM/test/matmul_vectorization.mlir
+++ b/iree/compiler/Conversion/LinalgToLLVM/test/matmul_vectorization.mlir
@@ -1,31 +1,31 @@
 // RUN: iree-opt --iree-codegen-linalg-to-llvm-matmul-vectorization-pass -split-input-file %s | IreeFileCheck %s
 
-// CHECK-LABEL: func @matmul_128x128x128
-// CHECK-SAME: (%[[ARG0:.+]]: memref<128x128xf32>, %[[ARG1:.+]]: memref<128x128xf32>, %[[ARG2:.+]]: memref<128x128xf32>)
 func @matmul_128x128x128(%arg0 : memref<128x128xf32>, %arg1: memref<128x128xf32>, %arg2: memref<128x128xf32>) {
     linalg.matmul %arg0, %arg1, %arg2 : (memref<128x128xf32>, memref<128x128xf32>, memref<128x128xf32>)
     return
 }
-// CHECK: %[[L3END:.+]] = constant 128 : index
-// CHECK: %[[L3STEP:.+]] = constant 64 : index
-// CHECK: %[[L1STEP:.+]] = constant 4 : index
-// CHECK: %[[L2STEP:.+]] = constant 32 : index
-// CHECK: %[[START:.+]] = constant 0 : index
-// CHECK: scf.for %[[IL3:.+]] = %[[START]] to %[[L3END]] step %[[L3STEP]]
-// CHECK: scf.for %[[JL3:.+]] = %[[START]] to %[[L3END]] step %[[L3STEP]]
-// CHECK: scf.for %[[KL3:.+]] = %[[START]] to %[[L3END]] step %[[L3STEP]]
-// CHECK: %[[ARG0_TILE_L3:.+]] = subview %[[ARG0]][%[[IL3]], %[[KL3]]] [64, 64] [1, 1] : memref<128x128xf32> to memref<64x64xf32
-// CHECK: %[[ARG1_TILE_L3:.+]] = subview %[[ARG1]][%[[KL3]], %[[JL3]]] [64, 64] [1, 1] : memref<128x128xf32> to memref<64x64xf32
-// CHECK: %[[ARG2_TILE_L3:.+]] = subview %[[ARG2]][%[[IL3]], %[[JL3]]] [64, 64] [1, 1] : memref<128x128xf32> to memref<64x64xf32
-// CHECK: scf.for %[[IL2:.+]] = %[[START]] to %[[L3STEP]] step %[[L2STEP]]
-// CHECK: scf.for %[[JL2:.+]] = %[[START]] to %[[L3STEP]] step %[[L2STEP]]
-// CHECK: scf.for %[[KL2:.+]] = %[[START]] to %[[L3STEP]] step %[[L2STEP]]
-// CHECK: %[[ARG0_TILE_L2:.+]] = subview %[[ARG0_TILE_L3]][%[[IL2]], %[[KL2]]] [32, 32] [1, 1] : memref<64x64xf32
-// CHECK: %[[ARG1_TILE_L2:.+]] = subview %[[ARG1_TILE_L3]][%[[KL2]], %[[JL2]]] [32, 32] [1, 1] : memref<64x64xf32
-// CHECK: %[[ARG2_TILE_L2:.+]] = subview %[[ARG2_TILE_L3]][%[[IL2]], %[[JL2]]] [32, 32] [1, 1] : memref<64x64xf32
-// CHECK: scf.for %[[IL1:.+]] = %[[START]] to %[[L2STEP]] step %[[L1STEP]]
-// CHECK: scf.for %[[JL1:.+]] = %[[START]] to %[[L2STEP]] step %[[L1STEP]]
-// CHECK: scf.for %[[KL1:.+]] = %[[START]] to %[[L2STEP]] step %[[L1STEP]]
-// CHECK: %[[ARG0_TILE_L1:.+]] = subview %[[ARG0_TILE_L2]][%[[IL1]], %[[KL1]]] [4, 4] [1, 1] : memref<32x32xf32
-// CHECK: %[[ARG1_TILE_L1:.+]] = subview %[[ARG1_TILE_L2]][%[[KL1]], %[[JL1]]] [4, 4] [1, 1] : memref<32x32xf32
-// CHECK: %[[ARG2_TILE_L1:.+]] = subview %[[ARG2_TILE_L2]][%[[IL1]], %[[JL1]]] [4, 4] [1, 1] : memref<32x32xf32
+// CHECK-LABEL: func @matmul_128x128x128
+//  CHECK-SAME: (%[[ARG0:.+]]: memref<128x128xf32>, %[[ARG1:.+]]: memref<128x128xf32>, %[[ARG2:.+]]: memref<128x128xf32>)
+//   CHECK-DAG: %[[L3END:.+]] = constant 128 : index
+//   CHECK-DAG: %[[L3STEP:.+]] = constant 64 : index
+//   CHECK-DAG: %[[L1STEP:.+]] = constant 4 : index
+//   CHECK-DAG: %[[L2STEP:.+]] = constant 32 : index
+//   CHECK-DAG: %[[START:.+]] = constant 0 : index
+//       CHECK: scf.for %[[IL3:.+]] = %[[START]] to %[[L3END]] step %[[L3STEP]]
+//       CHECK:   scf.for %[[JL3:.+]] = %[[START]] to %[[L3END]] step %[[L3STEP]]
+//       CHECK:     scf.for %[[KL3:.+]] = %[[START]] to %[[L3END]] step %[[L3STEP]]
+//       CHECK:       %[[ARG0_TILE_L3:.+]] = subview %[[ARG0]][%[[IL3]], %[[KL3]]] [64, 64] [1, 1] : memref<128x128xf32> to memref<64x64xf32
+//       CHECK:       %[[ARG1_TILE_L3:.+]] = subview %[[ARG1]][%[[KL3]], %[[JL3]]] [64, 64] [1, 1] : memref<128x128xf32> to memref<64x64xf32
+//       CHECK:       %[[ARG2_TILE_L3:.+]] = subview %[[ARG2]][%[[IL3]], %[[JL3]]] [64, 64] [1, 1] : memref<128x128xf32> to memref<64x64xf32
+//       CHECK:       scf.for %[[IL2:.+]] = %[[START]] to %[[L3STEP]] step %[[L2STEP]]
+//       CHECK:         scf.for %[[JL2:.+]] = %[[START]] to %[[L3STEP]] step %[[L2STEP]]
+//       CHECK:           scf.for %[[KL2:.+]] = %[[START]] to %[[L3STEP]] step %[[L2STEP]]
+//       CHECK:             %[[ARG0_TILE_L2:.+]] = subview %[[ARG0_TILE_L3]][%[[IL2]], %[[KL2]]] [32, 32] [1, 1] : memref<64x64xf32
+//       CHECK:             %[[ARG1_TILE_L2:.+]] = subview %[[ARG1_TILE_L3]][%[[KL2]], %[[JL2]]] [32, 32] [1, 1] : memref<64x64xf32
+//       CHECK:             %[[ARG2_TILE_L2:.+]] = subview %[[ARG2_TILE_L3]][%[[IL2]], %[[JL2]]] [32, 32] [1, 1] : memref<64x64xf32
+//       CHECK:             scf.for %[[IL1:.+]] = %[[START]] to %[[L2STEP]] step %[[L1STEP]]
+//       CHECK:               scf.for %[[JL1:.+]] = %[[START]] to %[[L2STEP]] step %[[L1STEP]]
+//       CHECK:                 scf.for %[[KL1:.+]] = %[[START]] to %[[L2STEP]] step %[[L1STEP]]
+//       CHECK:                   %[[ARG0_TILE_L1:.+]] = subview %[[ARG0_TILE_L2]][%[[IL1]], %[[KL1]]] [4, 4] [1, 1] : memref<32x32xf32
+//       CHECK:                   %[[ARG1_TILE_L1:.+]] = subview %[[ARG1_TILE_L2]][%[[KL1]], %[[JL1]]] [4, 4] [1, 1] : memref<32x32xf32
+//       CHECK:                   %[[ARG2_TILE_L1:.+]] = subview %[[ARG2_TILE_L2]][%[[IL1]], %[[JL1]]] [4, 4] [1, 1] : memref<32x32xf32
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/Attributes.h b/iree/compiler/Conversion/LinalgToSPIRV/Attributes.h
index e801c90..84d6393 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/Attributes.h
+++ b/iree/compiler/Conversion/LinalgToSPIRV/Attributes.h
@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef IREE_COMPILER_TRANSLATION_SPIRV_LINALGTOSPIRV_ATTRIBUTES_H_
-#define IREE_COMPILER_TRANSLATION_SPIRV_LINALGTOSPIRV_ATTRIBUTES_H_
+#ifndef IREE_COMPILER_CONVERSION_LINALGTOSPIRV_ATTRIBUTES_H_
+#define IREE_COMPILER_CONVERSION_LINALGTOSPIRV_ATTRIBUTES_H_
 
 #include "llvm/ADT/StringRef.h"
 
@@ -26,34 +26,13 @@
   return "vkspv.entry_point_schedule";
 }
 
-/// Enumerates the methods used to compute the number of workgroups to use with
-/// an entry point function. The lowering to SPIR-V sets an integer attribute on
-/// the entry point function with one of these values. It is later used by
-/// `recordDispatch` to compute the number of workgroups for the entry point
-/// function.
-enum class WorkgroupCountMethodology {
-  // TODO(#2134): Remove the `Default` option. This is only a fallback used for
-  // convolution/pooling cases that are currently not working as intended, as
-  // described in the bug.
-  Default = 0,               // Use the default mechanism used by IREE
-  LinearizeResultShape = 1,  // Use the linearized shape of the result of the
-                             // dispatch region
-  ResultShape = 2            // Use the shape of the dispatch region.
-};
-
-/// Returns the name of the attribute to use that propagates the method to use
-/// to compute the number of workgroups to use with an entry point function. The
-/// attribute used is an IntegerAttr with value being one of the enum entries of
-/// WorkgroupCountMethodology.
-// TODO(ravishankarm): The approach to use attributes to propagate the
-// methodology to use to compute number of workgroups is to convoluted. Ideally,
-// the lowering should produce a function that should then just be inlined at
-// the point this is needed.
-inline llvm::StringRef getWorkgroupCountAttrName() {
-  return "vkspv.workgroup_count_from_result_shape";
+/// Attribute on a entry point function that specifies which function computes
+/// the number of workgroups.
+inline llvm::StringRef getNumWorkgroupsFnAttrName() {
+  return "vkspv.num_workgroups_fn";
 }
 
 }  // namespace iree_compiler
 }  // namespace mlir
 
-#endif  // IREE_COMPILER_TRANSLATION_SPIRV_LINALGTOSPIRV_ATTRIBUTES_H_
+#endif  // IREE_COMPILER_CONVERSION_LINALGTOSPIRV_ATTRIBUTES_H_
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/BUILD b/iree/compiler/Conversion/LinalgToSPIRV/BUILD
index 6a30713..dd6717a 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/BUILD
+++ b/iree/compiler/Conversion/LinalgToSPIRV/BUILD
@@ -24,6 +24,9 @@
         "ConvertToGPUPass.cpp",
         "ConvertToSPIRVPass.cpp",
         "CooperativeMatrixAnalysis.cpp",
+        "DeclareNumWorkgroupsFnPass.cpp",
+        "KernelDispatchUtils.cpp",
+        "LegalizeNumWorkgroupsFnPass.cpp",
         "LinalgTileAndFusePass.cpp",
         "MarkerUtils.cpp",
         "MatMulVectorizationTest.cpp",
@@ -36,6 +39,7 @@
     hdrs = [
         "Attributes.h",
         "CooperativeMatrixAnalysis.h",
+        "KernelDispatchUtils.h",
         "MarkerUtils.h",
         "MemorySpace.h",
         "Passes.h",
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/CMakeLists.txt b/iree/compiler/Conversion/LinalgToSPIRV/CMakeLists.txt
index 8e34e87..ce9a6cb 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/CMakeLists.txt
+++ b/iree/compiler/Conversion/LinalgToSPIRV/CMakeLists.txt
@@ -20,6 +20,7 @@
   HDRS
     "Attributes.h"
     "CooperativeMatrixAnalysis.h"
+    "KernelDispatchUtils.h"
     "MarkerUtils.h"
     "MemorySpace.h"
     "Passes.h"
@@ -28,6 +29,9 @@
     "ConvertToGPUPass.cpp"
     "ConvertToSPIRVPass.cpp"
     "CooperativeMatrixAnalysis.cpp"
+    "DeclareNumWorkgroupsFnPass.cpp"
+    "KernelDispatchUtils.cpp"
+    "LegalizeNumWorkgroupsFnPass.cpp"
     "LinalgTileAndFusePass.cpp"
     "MarkerUtils.cpp"
     "MatMulVectorizationTest.cpp"
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/ConvertToGPUPass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/ConvertToGPUPass.cpp
index 38794f4..7f82188 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/ConvertToGPUPass.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/ConvertToGPUPass.cpp
@@ -21,11 +21,14 @@
 #include <array>
 #include <numeric>
 
+#include "iree/compiler/Conversion/CodegenUtils/FunctionUtils.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/Attributes.h"
+#include "iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/MemorySpace.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/Passes.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/Utils.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeDialect.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
@@ -506,15 +509,16 @@
 
 namespace {
 /// Pass to convert from tiled and fused linalg ops into gpu.func.
-struct ConvertToGPUPass : public PassWrapper<ConvertToGPUPass, FunctionPass> {
+struct ConvertToGPUPass
+    : public PassWrapper<ConvertToGPUPass, OperationPass<ModuleOp>> {
   ConvertToGPUPass() = default;
   ConvertToGPUPass(const ConvertToGPUPass &pass) {}
 
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<AffineDialect, gpu::GPUDialect, scf::SCFDialect>();
+    registry.insert<AffineDialect, gpu::GPUDialect, scf::SCFDialect,
+                    ShapeDialect>();
   }
-
-  void runOnFunction() override;
+  void runOnOperation() override;
 };
 
 struct SerializeParallelLoopPattern
@@ -651,11 +655,29 @@
         workgroupSize = {32, 1, 1};
       }
     }
-    rewriter.eraseOp(linalgOp);
     if (failed(updateWorkGroupSize(funcOp, workgroupSize))) return failure();
-    funcOp.setAttr(getWorkgroupCountAttrName(),
-                   rewriter.getI32IntegerAttr(static_cast<int32_t>(
-                       WorkgroupCountMethodology::LinearizeResultShape)));
+
+    // TODO(#3145): The use of the generated function to compute the number of
+    // workgroups works for dynamic shapes as well, but there is an issue with
+    // the shape computation + ha.device.switch creation that needs to be
+    // resolved before the generated function can be used on the host side. So
+    // disabling this approach for dynamic shape case.
+    if (llvm::all_of(linalgOp.getOperands(), [](Value v) {
+          Type t = v.getType();
+          return (t.isa<ShapedType>() &&
+                  t.cast<ShapedType>().hasStaticShape()) ||
+                 t.isIntOrIndexOrFloat();
+        })) {
+      if (funcOp.getAttr(getNumWorkgroupsFnAttrName()) &&
+          failed(createNumWorkgroupsFromLinearizedResultShape(
+              rewriter, cast<linalg::LinalgOp>(linalgOp.getOperation()), funcOp,
+              workgroupSize[0]))) {
+        return failure();
+      }
+    } else {
+      funcOp.removeAttr(getNumWorkgroupsFnAttrName());
+    }
+    rewriter.eraseOp(linalgOp);
     return success();
   }
 };
@@ -730,15 +752,7 @@
   patterns.insert<TileAndDistributeCopyOp>(context);
 }
 
-void ConvertToGPUPass::runOnFunction() {
-  FuncOp funcOp = getFunction();
-
-  Region &body = funcOp.getBody();
-  if (!llvm::hasSingleElement(body)) {
-    funcOp.emitError("unhandled dispatch function with multiple blocks");
-    return signalPassFailure();
-  }
-
+void ConvertToGPUPass::runOnOperation() {
   MLIRContext *context = &getContext();
   ConversionTarget target(*context);
   // After this pass Linalg and scf.parallel ops should be gone.
@@ -766,11 +780,19 @@
                   MapLinalgOpToLocalInvocationId<linalg::PoolingSumOp>,
                   RemoveLinalgRange, SerializeParallelLoopPattern>(context);
 
-  if (failed(applyFullConversion(funcOp, target, patterns)))
-    return signalPassFailure();
+  for (FuncOp funcOp : getOperation().getOps<FuncOp>()) {
+    if (!isEntryPoint(funcOp)) continue;
+    Region &body = funcOp.getBody();
+    if (!llvm::hasSingleElement(body)) {
+      funcOp.emitError("unhandled dispatch function with multiple blocks");
+      return signalPassFailure();
+    }
+    if (failed(applyFullConversion(funcOp, target, patterns)))
+      return signalPassFailure();
+  }
 }
 
-std::unique_ptr<OperationPass<FuncOp>> createConvertToGPUPass() {
+std::unique_ptr<OperationPass<ModuleOp>> createConvertToGPUPass() {
   return std::make_unique<ConvertToGPUPass>();
 }
 
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/DeclareNumWorkgroupsFnPass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/DeclareNumWorkgroupsFnPass.cpp
new file mode 100644
index 0000000..3bd8074
--- /dev/null
+++ b/iree/compiler/Conversion/LinalgToSPIRV/DeclareNumWorkgroupsFnPass.cpp
@@ -0,0 +1,152 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===- DeclareNumWorkgroupsFnPass.cpp - Declares num_workgroups_fn --------===//
+//
+// Define the function that computes the number of workgroups for every entry
+// point function. This pass only defines the function. Its body will be filled
+// in later.
+//
+//===----------------------------------------------------------------------===//
+#include "iree/compiler/Conversion/CodegenUtils/FunctionUtils.h"
+#include "iree/compiler/Conversion/LinalgToSPIRV/Attributes.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeDialect.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/StandardTypes.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+static constexpr const char kNumWorkgroupsStr[] = "__num_workgroups__";
+
+namespace {
+
+/// The contract between the host and the device is captured by the _impl
+/// function that is called from the main entry point function. This pattern
+/// looks for the call operation and
+/// - Declares (doesnt define) the function that computes the number of
+///   workgroups to use for this entry point function. It is defined later in
+///   the codegen pipeline, when the computation is mapped to
+///   workgroups/workitems. The signature of this function is
+///
+///      (!shapex.ranked_shape, !shapex.ranked_shape, ....) ->
+///      (index, index, index)
+///
+///   where the arguments are the shape of the tensor inputs + outputs of the
+///   dispatch region.
+/// - Sets the attribute `operand_result_index` on the
+///   `hal.interface.load.tensor`/`hal.interface.store.tensor` ops that are
+///   later used in the generation of the function declared here.
+struct DeclareNumWorkgroupsFn : OpRewritePattern<FuncOp> {
+  using OpRewritePattern<FuncOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(FuncOp entryPointFn,
+                                PatternRewriter &rewriter) const override {
+    if (!isEntryPoint(entryPointFn) ||
+        entryPointFn.getAttr(getNumWorkgroupsFnAttrName()))
+      return failure();
+    Region &body = entryPointFn.getBody();
+    if (!llvm::hasSingleElement(body)) {
+      return entryPointFn.emitError(
+          "unhandled dispatch function with multiple blocks");
+    }
+    auto callOps = body.front().getOps<CallOp>();
+    if (!llvm::hasSingleElement(callOps)) {
+      return entryPointFn.emitError(
+          "expected dispatch function to have a single call operation");
+    }
+    CallOp callOp = *callOps.begin();
+
+    SmallVector<ShapedType, 4> shapedTypes;
+    shapedTypes.reserve(callOp.getNumOperands() - 1 + callOp.getNumResults());
+
+    // Add `operand_result_index` attribute to `hal.interface.load.tensor`
+    // operations that define the operands of the call op.
+    for (Value operand : callOp.operands()) {
+      if (!operand.getType().isa<ShapedType>()) continue;
+      if (auto definingOp =
+              operand.getDefiningOp<IREE::HAL::InterfaceLoadTensorOp>()) {
+        definingOp.setAttr(getOperandResultNumAttrName(),
+                           rewriter.getI32IntegerAttr(shapedTypes.size()));
+      }
+      shapedTypes.push_back(operand.getType().cast<ShapedType>());
+    }
+
+    // Add `operand_result_index` attribute to the `hal.interface.store.tensor`
+    // that use the value returned by the call op.
+    for (Value result : callOp.getResults()) {
+      if (!result.getType().isa<ShapedType>()) continue;
+      for (auto &use : result.getUses()) {
+        if (auto storeOp =
+                dyn_cast<IREE::HAL::InterfaceStoreTensorOp>(use.getOwner())) {
+          storeOp.setAttr(getOperandResultNumAttrName(),
+                          rewriter.getI32IntegerAttr(shapedTypes.size()));
+        }
+      }
+      shapedTypes.push_back(result.getType().cast<ShapedType>());
+    }
+
+    IndexType indexType = rewriter.getIndexType();
+    SmallVector<Type, 4> argTypes = llvm::to_vector<4>(
+        llvm::map_range(shapedTypes, [&rewriter](ShapedType t) -> Type {
+          return Shape::RankedShapeType::get(t.getShape(),
+                                             rewriter.getContext());
+        }));
+    FuncOp numWorkgroupsFn = rewriter.create<FuncOp>(
+        entryPointFn.getLoc(), entryPointFn.getName().str() + kNumWorkgroupsStr,
+        rewriter.getFunctionType(argTypes, {indexType, indexType, indexType}));
+    numWorkgroupsFn.setVisibility(FuncOp::Visibility::Private);
+    entryPointFn.setAttr(getNumWorkgroupsFnAttrName(),
+                         rewriter.getSymbolRefAttr(numWorkgroupsFn));
+    rewriter.updateRootInPlace(entryPointFn, []() {});
+    return success();
+  }
+};
+
+/// Pass to define the function for number of workgroups for every entry point
+/// function.
+struct DeclareNumWorkgroupsFnPass
+    : public PassWrapper<DeclareNumWorkgroupsFnPass, OperationPass<ModuleOp>> {
+  DeclareNumWorkgroupsFnPass() = default;
+  DeclareNumWorkgroupsFnPass(const DeclareNumWorkgroupsFnPass &pass) {}
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<ShapeDialect>();
+  }
+};
+}  // namespace
+
+void DeclareNumWorkgroupsFnPass::runOnOperation() {
+  OwningRewritePatternList patterns;
+  MLIRContext *context = &getContext();
+  patterns.insert<DeclareNumWorkgroupsFn>(context);
+  applyPatternsAndFoldGreedily(getOperation(), patterns);
+}
+
+std::unique_ptr<OperationPass<ModuleOp>> createDeclareNumWorkgroupsFnPass() {
+  return std::make_unique<DeclareNumWorkgroupsFnPass>();
+}
+
+static PassRegistration<DeclareNumWorkgroupsFnPass> pass(
+    "iree-codegen-init-num-workgroups-fn",
+    "Declares the function that computes the number of workgroups to use",
+    [] { return std::make_unique<DeclareNumWorkgroupsFnPass>(); });
+
+}  // namespace iree_compiler
+}  // namespace mlir
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp
new file mode 100644
index 0000000..ea85ec2
--- /dev/null
+++ b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.cpp
@@ -0,0 +1,157 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//===- KernelDispatchUtils.cpp - Utilities for generating dispatch info ---===//
+//
+// This file defines utility functions that can be used to create information
+// the dispatch on the host side needs to execute an entry point function, like
+// the number of workgroups to use for launch, etc.
+//
+//===----------------------------------------------------------------------===//
+#include "iree/compiler/Conversion/CodegenUtils/FunctionUtils.h"
+#include "iree/compiler/Conversion/LinalgToSPIRV/Attributes.h"
+#include "iree/compiler/Dialect/IREE/IR/IREEOps.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/Module.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/StandardTypes.h"
+
+#define DEBUG_TYPE "kernel-dispatch-utils"
+
+namespace mlir {
+namespace iree_compiler {
+
+FuncOp getNumWorkgroupsFn(FuncOp entryPointFn) {
+  SymbolRefAttr attr =
+      entryPointFn.getAttrOfType<SymbolRefAttr>(getNumWorkgroupsFnAttrName());
+  if (!attr) {
+    entryPointFn.emitError("missing attribute '")
+        << getNumWorkgroupsFnAttrName() << "'";
+    return nullptr;
+  }
+  FuncOp numWorkgroupsFn = dyn_cast_or_null<FuncOp>(SymbolTable::lookupSymbolIn(
+      entryPointFn.getParentOfType<ModuleOp>(), attr));
+  if (!numWorkgroupsFn) {
+    entryPointFn.emitError("unable to find num workgroups fn ") << attr;
+    return nullptr;
+  }
+  if (!numWorkgroupsFn.empty()) {
+    entryPointFn.emitError("num workgroups fn expected to be empty");
+    return nullptr;
+  }
+  return numWorkgroupsFn;
+}
+
+/// Computes the bounds of the parallel loops partitioned across workgroups.
+static Optional<SmallVector<Value, 2>> getParallelLoopRange(
+    PatternRewriter &rewriter, Location loc, linalg::LinalgOp linalgOp) {
+  FuncOp numWorkgroupsFn =
+      getNumWorkgroupsFn(linalgOp.getParentOfType<FuncOp>());
+  if (!numWorkgroupsFn) return {};
+  LLVM_DEBUG({
+    llvm::dbgs() << "Found num workgroups function : "
+                 << numWorkgroupsFn.getName();
+  });
+  rewriter.setInsertionPointToEnd(numWorkgroupsFn.addEntryBlock());
+  llvm::SetVector<Operation *> slice;
+  getBackwardSlice(linalgOp, &slice);
+  BlockAndValueMapping mapper;
+  for (Operation *op : slice) {
+    rewriter.clone(*op, mapper);
+  }
+  // Clone the linalg operation just to compute the loop bounds.
+  linalg::LinalgOp clonedLinalgOp =
+      rewriter.clone(*linalgOp.getOperation(), mapper);
+  Optional<SmallVector<Value, 4>> bounds =
+      getLoopRanges(rewriter, clonedLinalgOp);
+  unsigned numParallelLoops = linalgOp.iterator_types()
+                                  .getValue()
+                                  .take_while([](Attribute attr) -> bool {
+                                    return attr.cast<StringAttr>().getValue() ==
+                                           getParallelIteratorTypeName();
+                                  })
+                                  .size();
+  SmallVector<Value, 2> returnVals(
+      bounds->begin(), std::next(bounds->begin(), numParallelLoops));
+  rewriter.eraseOp(clonedLinalgOp);
+  return returnVals;
+}
+
+/// Utility method to build IR that computes ceil(`numerator` / `denominator`)
+static Value buildCeilDiv(PatternRewriter &rewriter, Location loc,
+                          Value numerator, Value denominator) {
+  Value one = rewriter.create<ConstantIndexOp>(loc, 1);
+  Value t = rewriter.create<AddIOp>(
+      loc, numerator, rewriter.create<SubIOp>(loc, denominator, one));
+  return rewriter.create<SignedDivIOp>(loc, t, denominator);
+}
+
+/// Utility method to build IR that computes ceil(`numerator` / `denominator`)
+/// when denominator is a constant.
+static Value buildCeilDivConstDenominator(PatternRewriter &rewriter,
+                                          Location loc, Value numerator,
+                                          int64_t denominator) {
+  return buildCeilDiv(rewriter, loc, numerator,
+                      rewriter.create<ConstantIndexOp>(loc, denominator));
+}
+
+LogicalResult createNumWorkgroupsFromResultShape(PatternRewriter &rewriter,
+                                                 linalg::LinalgOp linalgOp,
+                                                 FuncOp entryPointFn,
+                                                 ArrayRef<int64_t> tileSizes) {
+  Location loc = linalgOp.getLoc();
+  OpBuilder::InsertionGuard gaurd(rewriter);
+  Optional<SmallVector<Value, 2>> parallelLoopRange =
+      getParallelLoopRange(rewriter, loc, linalgOp);
+  if (!parallelLoopRange) return failure();
+  Value one = rewriter.create<ConstantIndexOp>(loc, 1);
+  SmallVector<Value, 3> returnValues(3, one);
+  for (size_t i = 0, e = std::min<size_t>(parallelLoopRange->size(), 3); i != e;
+       ++i) {
+    returnValues[i] = buildCeilDivConstDenominator(
+        rewriter, loc, (*parallelLoopRange)[e - i - 1], tileSizes[e - i - 1]);
+  }
+  rewriter.create<mlir::ReturnOp>(loc, returnValues);
+  return success();
+}
+
+LogicalResult createNumWorkgroupsFromLinearizedResultShape(
+    PatternRewriter &rewriter, linalg::LinalgOp linalgOp, FuncOp entryPointFn,
+    int64_t workgroupSizeX) {
+  Location loc = linalgOp.getLoc();
+  OpBuilder::InsertionGuard gaurd(rewriter);
+  Optional<SmallVector<Value, 2>> parallelLoopRange =
+      getParallelLoopRange(rewriter, loc, linalgOp);
+  if (!parallelLoopRange) return failure();
+  Value one = rewriter.create<ConstantIndexOp>(loc, 1);
+  SmallVector<Value, 3> returnValues(3, one);
+  for (auto range : *parallelLoopRange) {
+    returnValues[0] = rewriter.create<MulIOp>(loc, range, returnValues[0]);
+  }
+  returnValues[0] = buildCeilDivConstDenominator(rewriter, loc, returnValues[0],
+                                                 workgroupSizeX);
+  rewriter.create<mlir::ReturnOp>(loc, returnValues);
+  return success();
+}
+
+}  // namespace iree_compiler
+}  // namespace mlir
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h
new file mode 100644
index 0000000..1573c55
--- /dev/null
+++ b/iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h
@@ -0,0 +1,66 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//===- KernelDispatchUtils.h - Utilities for generating dispatch info -----===//
+//
+// This file declares utility functions that can be used to create information
+// the dispatch on the host side needs to execute an entry point function, like
+// the number of workgroups to use for launch, etc.
+//
+//===----------------------------------------------------------------------===//
+#ifndef IREE_COMPILER_CONVERSION_LINALGTOSPIRV_KERNELDISPATCHUTILS_H_
+#define IREE_COMPILER_CONVERSION_LINALGTOSPIRV_KERNELDISPATCHUTILS_H_
+
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+class FuncOp;
+class LogicalResult;
+class PatternRewriter;
+class ShapedType;
+class Value;
+
+namespace linalg {
+class LinalgOp;
+}
+
+namespace iree_compiler {
+
+/// Generates a function that computes the number of workgroups as
+///  [ceil(`parallelLoopRange`[2] / `tileSizes`[2]),
+///   ceil(`parallelLoopRange`[1] / `tileSizes`[1]),
+///   ceil(`parallelLoopRange`[0] / `tileSizes`[0])]
+/// where `parallelLoopRange` is the ranges of the parallel loops of `linalgOp`
+/// distributed across workgroups.
+LogicalResult createNumWorkgroupsFromResultShape(PatternRewriter &rewriter,
+                                                 linalg::LinalgOp linalgOp,
+                                                 FuncOp entryPointFn,
+                                                 ArrayRef<int64_t> tileSizes);
+
+/// Generates a function that computes the number of workgroups as
+///  ceil(`parallelLoopRange`[0] * `parallelLoopRange`[1] * ... *
+///       `parallelLoopRange`[n-1]  /  `workgroupSizeX`)
+/// where `parallelLoopRange` is the ranges of the parallel loops of `linalgOp`
+/// distributed across workgroups.
+LogicalResult createNumWorkgroupsFromLinearizedResultShape(
+    PatternRewriter &rewriter, linalg::LinalgOp linalgOp, FuncOp entryPointFn,
+    int64_t workgroupSizeX);
+
+/// For a given `entryPointFn` return the function that computes the number of
+/// workgroups to use at launch time.
+FuncOp getNumWorkgroupsFn(FuncOp entryPointFn);
+
+}  // namespace iree_compiler
+}  // namespace mlir
+#endif  // IREE_COMPILER_CONVERSION_LINALGTOSPIRV_DISPATCHUTILS_H_
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/LegalizeNumWorkgroupsFnPass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/LegalizeNumWorkgroupsFnPass.cpp
new file mode 100644
index 0000000..1d832c1
--- /dev/null
+++ b/iree/compiler/Conversion/LinalgToSPIRV/LegalizeNumWorkgroupsFnPass.cpp
@@ -0,0 +1,126 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//===-LegalizeNumWorkgroupsFnPass.cpp - Legalize to be runnable on host ---===//
+//
+// The function generated by the codegeneration pass to compute the number of
+// workgroups uses a slice of the device-side code. Legalize it to run on the
+// host.
+//
+//===----------------------------------------------------------------------===//
+#include "iree/compiler/Conversion/CodegenUtils/FunctionUtils.h"
+#include "iree/compiler/Conversion/LinalgToSPIRV/Attributes.h"
+#include "iree/compiler/Dialect/IREE/IR/IREEOps.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
+#include "mlir/IR/Module.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+namespace {
+
+/// Pattern to legalize shapex.tie_shape operation to tie the shape of the
+/// `iree.placeholder` result to the argument of the function.
+struct LegalizeTieShapeOp : OpRewritePattern<Shape::TieShapeOp> {
+  using OpRewritePattern<Shape::TieShapeOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(Shape::TieShapeOp tieShapeOp,
+                                PatternRewriter &rewriter) const override {
+    if (tieShapeOp.shape().isa<BlockArgument>()) return failure();
+    auto phOp = dyn_cast_or_null<IREE::PlaceholderOp>(
+        tieShapeOp.operand().getDefiningOp());
+    if (!phOp) return failure();
+    IntegerAttr operandNumAttr =
+        phOp.getAttrOfType<IntegerAttr>(getOperandResultNumAttrName());
+    if (!operandNumAttr) {
+      return phOp.emitRemark("expected operand_result_index attribute");
+    }
+    FuncOp numWorkgroupsFn = phOp.getParentOfType<FuncOp>();
+    rewriter.replaceOpWithNewOp<Shape::TieShapeOp>(
+        tieShapeOp, phOp,
+        numWorkgroupsFn.getArgument(
+            phOp.getAttrOfType<IntegerAttr>(getOperandResultNumAttrName())
+                .getInt()));
+    return success();
+  }
+};
+
+/// Pattern to remove dead `iree.placeholder` ops. They arent removed since they
+/// are tagged as having `MemoryEffect`.
+struct RemoveDeadPlaceholderOp : OpRewritePattern<IREE::PlaceholderOp> {
+  using OpRewritePattern<IREE::PlaceholderOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(IREE::PlaceholderOp phOp,
+                                PatternRewriter &rewriter) const override {
+    if (phOp.use_empty()) {
+      rewriter.eraseOp(phOp);
+      return success();
+    }
+    return failure();
+  }
+};
+
+/// Pass to legalize the function that computes the number of workgroups to use
+/// for launch to be runnable on the host-side.
+struct LegalizeNumWorkgroupsFnPass
+    : public PassWrapper<LegalizeNumWorkgroupsFnPass, OperationPass<ModuleOp>> {
+  LegalizeNumWorkgroupsFnPass() = default;
+  LegalizeNumWorkgroupsFnPass(const LegalizeNumWorkgroupsFnPass &pass) {}
+  void runOnOperation() override;
+};
+}  // namespace
+
+static void populateLegalizeNumWorkgroupsFnPattern(
+    MLIRContext *context, OwningRewritePatternList &patterns) {
+  patterns.insert<LegalizeTieShapeOp, RemoveDeadPlaceholderOp>(context);
+}
+
+void LegalizeNumWorkgroupsFnPass::runOnOperation() {
+  ModuleOp module = getOperation();
+  auto fns = module.getOps<FuncOp>();
+
+  OwningRewritePatternList patterns;
+  MLIRContext *context = &getContext();
+  populateLegalizeNumWorkgroupsFnPattern(context, patterns);
+
+  SymbolTable symbolTable(module.getOperation());
+  for (FuncOp fn : fns) {
+    if (!isEntryPoint(fn)) continue;
+    auto numWorkgroupsFnAttr =
+        fn.getAttrOfType<SymbolRefAttr>(getNumWorkgroupsFnAttrName());
+    if (!numWorkgroupsFnAttr) continue;
+    StringRef numWorkgroupsFnName = numWorkgroupsFnAttr.getLeafReference();
+    FuncOp numWorkgroupsFn = symbolTable.lookup<FuncOp>(numWorkgroupsFnName);
+    if (!numWorkgroupsFn) {
+      fn.emitError("unable to find function to compute number of workgroups ")
+          << numWorkgroupsFnName;
+      return signalPassFailure();
+    }
+    if (failed(applyPatternsAndFoldGreedily(numWorkgroupsFn, patterns)))
+      return signalPassFailure();
+  }
+}
+
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeNumWorkgroupsFnPass() {
+  return std::make_unique<LegalizeNumWorkgroupsFnPass>();
+}
+
+static PassRegistration<LegalizeNumWorkgroupsFnPass> pass(
+    "iree-codegen-legalize-num-workgroups-fn",
+    "Legalize the function that computes the number of workgroups to use to be "
+    "usable on the host side",
+    [] { return std::make_unique<LegalizeNumWorkgroupsFnPass>(); });
+
+}  // namespace iree_compiler
+}  // namespace mlir
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
index 71f550f..4c7417e 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/LinalgTileAndFusePass.cpp
@@ -17,11 +17,14 @@
 // Implements a pass to tile and fuse linalg operations on buffers.
 //
 //===----------------------------------------------------------------------===//
+#include "iree/compiler/Conversion/CodegenUtils/FunctionUtils.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/Attributes.h"
+#include "iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/MarkerUtils.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/MemorySpace.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/Passes.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/Utils.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeDialect.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
@@ -195,7 +198,7 @@
 namespace {
 /// Function pass that implements tiling and fusion in Linalg on buffers.
 struct LinalgTileAndFusePass
-    : public PassWrapper<LinalgTileAndFusePass, FunctionPass> {
+    : public PassWrapper<LinalgTileAndFusePass, OperationPass<ModuleOp>> {
   LinalgTileAndFusePass(ArrayRef<int64_t> workgroupSize = {},
                         ArrayRef<int64_t> tileSizes = {},
                         bool useWorkgroupMem = false) {
@@ -206,15 +209,10 @@
   LinalgTileAndFusePass(const LinalgTileAndFusePass &pass) {}
 
   void getDependentDialects(DialectRegistry &registry) const override {
-    // clang-format off
-    registry.insert<AffineDialect,
-                    gpu::GPUDialect,
-                    linalg::LinalgDialect,
-                    scf::SCFDialect>();
-    // clang-format on
+    registry.insert<AffineDialect, gpu::GPUDialect, linalg::LinalgDialect,
+                    scf::SCFDialect, ShapeDialect>();
   }
-
-  void runOnFunction() override;
+  void runOnOperation() override;
 
  private:
   Option<bool> useWorkgroupMemory{
@@ -243,17 +241,20 @@
 
 /// Pattern for tiling operations. Updates the workgroup size in the surrounding
 /// function operation if tiling succeeds.
-struct TileMatmulPattern
-    : public linalg::LinalgTilingPattern<linalg::MatmulOp> {
-  using Base = linalg::LinalgTilingPattern<linalg::MatmulOp>;
+template <typename MatmulOp>
+struct TileMatmulPattern : public linalg::LinalgBaseTilingPattern {
+  using Base = linalg::LinalgBaseTilingPattern;
   TileMatmulPattern(MLIRContext *context, linalg::LinalgTilingOptions options,
+                    ArrayRef<int64_t> tileSizes,
                     ArrayRef<int64_t> workgroupSize, PatternBenefit benefit = 1)
-      : Base(context, options.setDistributionOptions(matmulDistributionOptions),
+      : Base(MatmulOp::getOperationName(), context,
+             options.setDistributionOptions(matmulDistributionOptions),
              linalg::LinalgMarker(
                  ArrayRef<Identifier>(),
                  Identifier::get(getWorkgroupNumItemsGENumItersMarker(),
                                  context)),
              benefit),
+        tileSizes(tileSizes.begin(), tileSizes.end()),
         workgroupSize(workgroupSize.begin(), workgroupSize.end()) {}
 
   virtual LogicalResult matchAndRewrite(Operation *op,
@@ -262,44 +263,17 @@
     // erased.
     FuncOp funcOp = op->getParentOfType<FuncOp>();
     if (!funcOp || failed(Base::matchAndRewrite(op, rewriter)) ||
-        failed(updateWorkGroupSize(funcOp, workgroupSize)))
-      return failure();
-    funcOp.setAttr(getWorkgroupCountAttrName(),
-                   rewriter.getI32IntegerAttr(static_cast<int32_t>(
-                       WorkgroupCountMethodology::ResultShape)));
-    return success();
-  }
-
-  SmallVector<int64_t, 3> workgroupSize;
-};
-
-struct TileBatchMatmulPattern
-    : public linalg::LinalgTilingPattern<linalg::BatchMatmulOp> {
-  using Base = linalg::LinalgTilingPattern<linalg::BatchMatmulOp>;
-  TileBatchMatmulPattern(MLIRContext *context,
-                         linalg::LinalgTilingOptions options,
-                         ArrayRef<int64_t> workgroupSize,
-                         PatternBenefit benefit = 1)
-      : Base(context, options.setDistributionOptions(matmulDistributionOptions),
-             linalg::LinalgMarker(
-                 ArrayRef<Identifier>(),
-                 Identifier::get(getWorkgroupMarker(), context)),
-             benefit),
-        workgroupSize(workgroupSize.begin(), workgroupSize.end()) {}
-
-  LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override {
-    FuncOp funcOp = op->getParentOfType<FuncOp>();
-    if (!funcOp || failed(Base::matchAndRewrite(op, rewriter)) ||
-        failed(updateWorkGroupSize(funcOp, this->workgroupSize))) {
+        failed(updateWorkGroupSize(funcOp, workgroupSize)) ||
+        (funcOp.getAttr(getNumWorkgroupsFnAttrName()) &&
+         failed(createNumWorkgroupsFromResultShape(
+             rewriter, cast<linalg::LinalgOp>(op), funcOp, tileSizes)))) {
       return failure();
     }
-    funcOp.setAttr(getWorkgroupCountAttrName(),
-                   rewriter.getI32IntegerAttr(static_cast<int32_t>(
-                       WorkgroupCountMethodology::ResultShape)));
+    rewriter.eraseOp(op);
     return success();
   }
 
+  SmallVector<int64_t, 3> tileSizes;
   SmallVector<int64_t, 3> workgroupSize;
 };
 
@@ -318,6 +292,7 @@
 struct TileConvPoolPattern : public linalg::LinalgTilingPattern<OpTy> {
   using Base = linalg::LinalgTilingPattern<OpTy>;
   TileConvPoolPattern(MLIRContext *context, linalg::LinalgTilingOptions options,
+                      ArrayRef<int64_t> tileSizes,
                       ArrayRef<int64_t> workgroupSize,
                       PatternBenefit benefit = 1)
       : Base(context,
@@ -326,6 +301,7 @@
                  ArrayRef<Identifier>(),
                  Identifier::get(getWorkgroupMarker(), context)),
              benefit),
+        tileSizes(tileSizes.begin(), tileSizes.end()),
         workgroupSize(workgroupSize.begin(), workgroupSize.end()) {}
 
   LogicalResult matchAndRewrite(Operation *op,
@@ -335,12 +311,11 @@
     if (!funcOp || failed(Base::matchAndRewrite(op, rewriter)) ||
         failed(updateWorkGroupSize(funcOp, this->workgroupSize)))
       return failure();
-    funcOp.setAttr(getWorkgroupCountAttrName(),
-                   rewriter.getI32IntegerAttr(static_cast<int32_t>(
-                       WorkgroupCountMethodology::Default)));
+    funcOp.removeAttr(getNumWorkgroupsFnAttrName());
     return success();
   }
 
+  SmallVector<int64_t, 3> tileSizes;
   SmallVector<int64_t, 3> workgroupSize;
 };
 
@@ -406,66 +381,75 @@
 };
 }  // namespace
 
-void LinalgTileAndFusePass::runOnFunction() {
+void LinalgTileAndFusePass::runOnOperation() {
   MLIRContext *context = &getContext();
-  FuncOp funcOp = getFunction();
-  Region &body = funcOp.getBody();
-  if (!llvm::hasSingleElement(body.getBlocks())) {
-    funcOp.emitError("unhandled dispatch function with multiple blocks");
-    return signalPassFailure();
-  }
-  Block &block = body.front();
-  auto linalgOps = block.getOps<linalg::LinalgOp>();
-  if (linalgOps.empty()) return;
+  ModuleOp module = getOperation();
 
-  TileSizeCalculator tileSizeCalculator(funcOp);
-  if (tileSizes.empty()) {
-    // Get the tile sizes to use for the lowering.
-    SmallVector<int64_t, 3> tileSizes;
-    SmallVector<linalg::LinalgOp, 1> opsVec(linalgOps.begin(), linalgOps.end());
-    if (failed(tileSizeCalculator.inferTileAndWorkgroupSize(opsVec)))
+  LLVM_DEBUG(
+      llvm::dbgs() << "--- IREE Linalg tile and fuse configuration ---\n";);
+  for (FuncOp funcOp : module.getOps<FuncOp>()) {
+    if (!isEntryPoint(funcOp)) continue;
+
+    Region &body = funcOp.getBody();
+    if (!llvm::hasSingleElement(body.getBlocks())) {
+      funcOp.emitError("unhandled dispatch function with multiple blocks");
       return signalPassFailure();
-  } else {
-    tileSizeCalculator.setTileSizes(tileSizes);
-    if (!workgroupSize.empty())
-      tileSizeCalculator.setWorkgroupSize(workgroupSize);
-  }
+    }
+    Block &block = body.front();
+    auto linalgOps = block.getOps<linalg::LinalgOp>();
+    if (linalgOps.empty()) return;
 
-  LLVM_DEBUG({
-    llvm::dbgs() << "--- IREE Linalg tile and fuse configuration ---\n";
-    llvm::dbgs() << "# workgroup sizes: [";
-    interleaveComma(tileSizeCalculator.getWorkgroupSize(), llvm::dbgs());
-    llvm::dbgs() << "]\ntile sizes: [";
-    interleaveComma(tileSizeCalculator.getTileSizes(), llvm::dbgs());
-    llvm::dbgs() << "]\n";
-  });
+    TileSizeCalculator tileSizeCalculator(funcOp);
+    if (tileSizes.empty()) {
+      // Get the tile sizes to use for the lowering.
+      SmallVector<int64_t, 3> tileSizes;
+      SmallVector<linalg::LinalgOp, 1> opsVec(linalgOps.begin(),
+                                              linalgOps.end());
+      if (failed(tileSizeCalculator.inferTileAndWorkgroupSize(opsVec)))
+        return signalPassFailure();
+    } else {
+      tileSizeCalculator.setTileSizes(tileSizes);
+      if (!workgroupSize.empty())
+        tileSizeCalculator.setWorkgroupSize(workgroupSize);
+    }
 
-  OwningRewritePatternList tilingPatterns;
-  tilingPatterns
-      .insert<TileConvPoolPattern<linalg::ConvOp>, TileMatmulPattern,
-              TileBatchMatmulPattern, TileConvPoolPattern<linalg::PoolingMaxOp>,
-              TileConvPoolPattern<linalg::PoolingMinOp>,
-              TileConvPoolPattern<linalg::PoolingSumOp>>(
-          context,
-          linalg::LinalgTilingOptions()
-              .setTileSizes(tileSizeCalculator.getTileSizes())
-              .setLoopType(linalg::LinalgTilingLoopType::ParallelLoops),
-          tileSizeCalculator.getWorkgroupSize());
-  applyPatternsAndFoldGreedily(getOperation(), tilingPatterns);
+    LLVM_DEBUG({
+      llvm::dbgs() << "@func " << funcOp.getName() << ": # workgroup sizes: [";
+      interleaveComma(tileSizeCalculator.getWorkgroupSize(), llvm::dbgs());
+      llvm::dbgs() << "]\ntile sizes: [";
+      interleaveComma(tileSizeCalculator.getTileSizes(), llvm::dbgs());
+      llvm::dbgs() << "]\n";
+    });
 
-  if (useWorkgroupMemory) {
-    // The promotion patterns are put separate from the tiling patterns to make
-    // sure that the allocated scratchspace memory is constant sizes which
-    // requires some folding to trigger.
-    OwningRewritePatternList promotionPatterns;
-    promotionPatterns.insert<PromoteMatmulSubviewsPattern,
-                             PromoteConvolutionSubviewsPattern>(
+    OwningRewritePatternList tilingPatterns;
+    tilingPatterns.insert<TileConvPoolPattern<linalg::ConvOp>,
+                          TileMatmulPattern<linalg::MatmulOp>,
+                          TileMatmulPattern<linalg::BatchMatmulOp>,
+                          TileConvPoolPattern<linalg::PoolingMaxOp>,
+                          TileConvPoolPattern<linalg::PoolingMinOp>,
+                          TileConvPoolPattern<linalg::PoolingSumOp>>(
         context,
-        linalg::LinalgPromotionOptions()
-            .setAllocationDeallocationFns(allocateWorkgroupMemory,
-                                          deallocateWorkgroupMemory)
-            .setCopyInOutFns(copyToWorkgroupMemory, copyToWorkgroupMemory));
-    applyPatternsAndFoldGreedily(getOperation(), promotionPatterns);
+        linalg::LinalgTilingOptions()
+            .setTileSizes(tileSizeCalculator.getTileSizes())
+            .setLoopType(linalg::LinalgTilingLoopType::ParallelLoops),
+        tileSizeCalculator.getTileSizes(),
+        tileSizeCalculator.getWorkgroupSize());
+    applyPatternsAndFoldGreedily(funcOp, tilingPatterns);
+
+    if (useWorkgroupMemory) {
+      // The promotion patterns are put separate from the tiling patterns to
+      // make sure that the allocated scratchspace memory is constant sizes
+      // which requires some folding to trigger.
+      OwningRewritePatternList promotionPatterns;
+      promotionPatterns.insert<PromoteMatmulSubviewsPattern,
+                               PromoteConvolutionSubviewsPattern>(
+          context,
+          linalg::LinalgPromotionOptions()
+              .setAllocationDeallocationFns(allocateWorkgroupMemory,
+                                            deallocateWorkgroupMemory)
+              .setCopyInOutFns(copyToWorkgroupMemory, copyToWorkgroupMemory));
+      applyPatternsAndFoldGreedily(funcOp, promotionPatterns);
+    }
   }
 }
 
@@ -473,7 +457,7 @@
 // Pass entry point and registration
 //===----------------------------------------------------------------------===//
 
-std::unique_ptr<OperationPass<FuncOp>> createLinalgTileAndFusePass(
+std::unique_ptr<OperationPass<ModuleOp>> createLinalgTileAndFusePass(
     ArrayRef<int64_t> workgroupSize, ArrayRef<int64_t> tileSizes,
     bool useWorkgroupMemory) {
   return std::make_unique<LinalgTileAndFusePass>(workgroupSize, tileSizes,
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/Passes.cpp b/iree/compiler/Conversion/LinalgToSPIRV/Passes.cpp
index 5c1b672..cc90ccc 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/Passes.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/Passes.cpp
@@ -100,11 +100,9 @@
   //   afterwards. This gives each Linalg op a second chance to be tiled,
   //   with the second tile and fuse pass.
   //===--------------------------------------------------------------------===//
+  pm.addPass(createSplitDispatchFunctionPass());
   pm.addPass(createLinalgTileAndFusePass(
       options.workgroupSize, options.tileSizes, options.useWorkgroupMemory));
-  pm.addPass(createSplitDispatchFunctionPass());
-  pm.addPass(createLinalgTileAndFusePass(options.workgroupSize,
-                                         options.useWorkgroupMemory));
   pm.addPass(createCanonicalizerPass());
 
   //===--------------------------------------------------------------------===//
@@ -121,6 +119,16 @@
   pm.addPass(createCSEPass());
 
   //===--------------------------------------------------------------------===//
+  // Legalize the function that computes the number of workgroups to be runnable
+  // on the host.
+  //
+  // Post-conditions:
+  //   - The shape of the values created from `iree.placeholder` operations are
+  //     tied to the arguments of the function.
+  //===--------------------------------------------------------------------===//
+  pm.addPass(createLegalizeNumWorkgroupsFnPass());
+
+  //===--------------------------------------------------------------------===//
   // Resolve shape related ops.
   //
   // Pre-conditions:
@@ -137,6 +145,16 @@
   pm.addPass(createResolveShapeOpsPass());
 
   //===--------------------------------------------------------------------===//
+  // Legalize the function that computes the number of workgroups to be runnable
+  // on the host.
+  //
+  // Post-conditions:
+  //   - The dead `iree.placeholder` operations are removed after shape
+  //     resolution.
+  //===--------------------------------------------------------------------===//
+  pm.addPass(createLegalizeNumWorkgroupsFnPass());
+
+  //===--------------------------------------------------------------------===//
   // Prepare stdandard ops for SPIR-V conversion.
   //
   // Post-conditions:
@@ -174,6 +192,21 @@
 void buildSPIRVTransformPassPipeline(OpPassManager &pm,
                                      const SPIRVCodegenOptions &options) {
   //===--------------------------------------------------------------------===//
+  // The entry point functions call an _impl function that captures the ABI that
+  // the host side uses for the dispatch region. This ABI is needed when
+  // generating the function that computes the number of workgroups. Declare the
+  // function that returns the number of workgroups needed for an entry point
+  // function.
+  //
+  // Post-conditions
+
+  //   - An empty, private function is defined for each entry point function
+  //     that returns the number of workgroups.
+  //   - The entry point function gets an attribute `vkspv.num_workgroups_fn` to
+  //     record which function in the module returns the number of workgroups.
+  pm.addPass(createDeclareNumWorkgroupsFnPass());
+
+  //===--------------------------------------------------------------------===//
   // Inline the impl dispatch function into the wrapper dispatch function.
   //
   // TODO(antiagainst): re-evaluate the inlining timing.
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/Passes.h b/iree/compiler/Conversion/LinalgToSPIRV/Passes.h
index 72872f5..6e07cb0 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/Passes.h
+++ b/iree/compiler/Conversion/LinalgToSPIRV/Passes.h
@@ -29,19 +29,27 @@
   bool useWorkgroupMemory = false;
 };
 
+/// Pass to initialize the function that computes the number of workgroups for
+/// each entry point function. The function is defined, but is populated later.
+std::unique_ptr<OperationPass<ModuleOp>> createDeclareNumWorkgroupsFnPass();
+
 /// Pass to tile and fuse linalg operations on buffers. The pass takes as
 /// argument the `workgroupSize` that the tiling should use. Note that the
 /// tile-sizes are the reverse of the workgroup size. So workgroup size along
 /// "x" is used to tile the innermost loop, along "y" for the next innermost (if
 /// it exists) and along "z" for the next loop (if it exists). The workgroup
 /// size is expected to be of size at-most 3.
-std::unique_ptr<OperationPass<FuncOp>> createLinalgTileAndFusePass(
+std::unique_ptr<OperationPass<ModuleOp>> createLinalgTileAndFusePass(
     ArrayRef<int64_t> workGroupSize = {}, ArrayRef<int64_t> tileSizes = {},
     bool useWorkgroupMem = false);
 
 /// Pass to add the synchronizations and attributes needed to lower from PLoops
 /// to GPU dialect.
-std::unique_ptr<OperationPass<FuncOp>> createConvertToGPUPass();
+std::unique_ptr<OperationPass<ModuleOp>> createConvertToGPUPass();
+
+/// Pass to legalize function that returns number of workgroups to use for
+/// launch to be runnable on the host.
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeNumWorkgroupsFnPass();
 
 /// Pass to perform the final conversion to SPIR-V dialect.
 /// This pass converts remaining interface ops into SPIR-V global variables,
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/SplitDispatchFunctionPass.cpp b/iree/compiler/Conversion/LinalgToSPIRV/SplitDispatchFunctionPass.cpp
index 9a98f6f..94436c6 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/SplitDispatchFunctionPass.cpp
+++ b/iree/compiler/Conversion/LinalgToSPIRV/SplitDispatchFunctionPass.cpp
@@ -28,6 +28,7 @@
 
 #include "iree/compiler/Conversion/CodegenUtils/FunctionUtils.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/Attributes.h"
+#include "iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h"
 #include "iree/compiler/Conversion/LinalgToSPIRV/Passes.h"
 #include "iree/compiler/Dialect/IREE/IR/IREEOps.h"
 #include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
@@ -45,6 +46,8 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/RegionUtils.h"
 
+#define DEBUG_TYPE "split-dispatch-function"
+
 namespace mlir {
 namespace iree_compiler {
 
@@ -178,13 +181,32 @@
     StringRef newFnName = splitKernels.back();
     builder.setInsertionPointToStart(moduleOp.getBody());
     auto newFn = builder.create<FuncOp>(loc, newFnName, oldFn.getType());
+    LLVM_DEBUG({
+      llvm::dbgs() << "Created new function : func @" << newFn.getName()
+                   << "\n";
+    });
 
     // Copy over all attributes except type and name.
     for (const auto &namedAttr : oldFn.getAttrs()) {
       if (namedAttr.first != impl::getTypeAttrName() &&
-          namedAttr.first != SymbolTable::getSymbolAttrName())
+          namedAttr.first != SymbolTable::getSymbolAttrName() &&
+          namedAttr.first != getNumWorkgroupsFnAttrName())
         newFn.setAttr(namedAttr.first, namedAttr.second);
     }
+    // Need special handling for the number of workgroups function.
+    if (FuncOp numWorkgroupsFn = getNumWorkgroupsFn(oldFn)) {
+      FuncOp newNumWorkgroupsFn =
+          builder.create<FuncOp>(loc, newFnName.str() + "__num_workgroups__",
+                                 numWorkgroupsFn.getType());
+      newNumWorkgroupsFn.setVisibility(FuncOp::Visibility::Private);
+      newFn.setAttr(getNumWorkgroupsFnAttrName(),
+                    builder.getSymbolRefAttr(newNumWorkgroupsFn));
+      LLVM_DEBUG({
+        llvm::dbgs() << "Added func @" << newNumWorkgroupsFn.getName()
+                     << " as num workgroups fn for func @" << newFn.getName()
+                     << "\n";
+      });
+    }
 
     // Collect the closure for the current Linalg op.
     closure.clear();
@@ -211,6 +233,15 @@
   moduleOp.setAttr(getEntryPointScheduleAttrName(),
                    builder.getArrayAttr(entryPoints));
 
+  if (FuncOp numWorkgroupsFn = getNumWorkgroupsFn(oldFn)) {
+    LLVM_DEBUG({
+      llvm::dbgs() << "Erased num workgroups fn func @"
+                   << numWorkgroupsFn.getName() << " for func @"
+                   << oldFn.getName() << "\n";
+    });
+    numWorkgroupsFn.erase();
+  }
+  LLVM_DEBUG({ llvm::dbgs() << "Erased func @" << oldFn.getName() << "\n"; });
   oldFn.erase();
   return success();
 }
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu.mlir
index 51b088b..daa92e2 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu.mlir
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/convert_to_gpu.mlir
@@ -7,10 +7,13 @@
     [Shader], [SPV_KHR_storage_buffer_storage_class]>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @parallel_4D(%arg0: memref<?x?x?x?xf32>,
-                    %arg1 : memref<?x?x?x?xf32>,
-                    %arg2 : memref<?x?x?x?xf32>)
-  attributes {iree.dispatch_fn_name = "parallel_4D"} {
+  func @parallel_4D() {
+    %arg0 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg0, operand_result_index = 4 : i32} : memref<?x?x?x?xf32>
+    %arg1 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg1, operand_result_index = 9 : i32} : memref<?x?x?x?xf32>
+    %arg2 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@ret0, operand_result_index = 10 : i32} : memref<?x?x?x?xf32>
     linalg.generic
       {args_in = 2 : i64, args_out = 1 : i64,
        indexing_maps = [#map0, #map0, #map0],
@@ -22,10 +25,18 @@
     } : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
     return
   }
+  func @parallel_4D__num_workgroups__
+    (!shapex.ranked_shape<[?,?,?,?]>, !shapex.ranked_shape<[?,?,?,?]>,
+     !shapex.ranked_shape<[?,?,?,?]>) -> (index, index, index)
+    attributes {sym_visibility = "private"}
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
 // CHECK-LABEL: func @parallel_4D
 //  CHECK-SAME:   local_size = dense<[32, 1, 1]>
-//  CHECK-SAME:   vkspv.workgroup_count_from_result_shape = 1
 //   CHECK-DAG:     %[[C0:.+]] = constant 0 : index
 //   CHECK-DAG:     %[[C1:.+]] = constant 1 : index
 //   CHECK-DAG:     %[[C2:.+]] = constant 2 : index
@@ -56,6 +67,71 @@
 
 // -----
 
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+module attributes {
+  spv.target_env =
+    #spv.target_env<#spv.vce<v1.3,
+    [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
+  func @parallel_4D_static() attributes {vkspv.num_workgroups_fn = @parallel_4D_static__num_workgroups__} {
+    %arg0 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<3x4x5x6xf32>
+    %arg1 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg1, operand_result_index = 1 : i32} : memref<3x4x5x6xf32>
+    %arg2 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@ret0, operand_result_index = 2 : i32} : memref<3x4x5x6xf32>
+    linalg.generic
+      {args_in = 2 : i64, args_out = 1 : i64,
+       indexing_maps = [#map0, #map0, #map0],
+       iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+      %arg0, %arg1, %arg2 {
+    ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
+      %0 = addf %arg3, %arg4 : f32
+      linalg.yield %0 : f32
+    } : memref<3x4x5x6xf32>, memref<3x4x5x6xf32>, memref<3x4x5x6xf32>
+    return
+  }
+  func @parallel_4D_static__num_workgroups__
+    (!shapex.ranked_shape<[3,4,5,6]>, !shapex.ranked_shape<[3,4,5,6]>,
+     !shapex.ranked_shape<[3,4,5,6]>) -> (index, index, index)
+    attributes {sym_visibility = "private"}
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
+}
+// CHECK-LABEL: func @parallel_4D_static()
+//  CHECK-SAME:   local_size = dense<[32, 1, 1]>
+//  CHECK-SAME:   vkspv.num_workgroups_fn = @[[NUM_WORKGROUPS_FN:[a-zA-Z0-9_]+]]
+//   CHECK-DAG:     %[[C360:.+]] = constant 360 : index
+//   CHECK-DAG:     %[[C120:.+]] = constant 120 : index
+//   CHECK-DAG:     %[[C30:.+]] = constant 30 : index
+//   CHECK-DAG:     %[[C6:.+]] = constant 6 : index
+//   CHECK-DAG:     %[[BID:.+]] = "gpu.block_id"() {dimension = "x"}
+//   CHECK-DAG:     %[[BDIM:.+]] = "gpu.block_dim"() {dimension = "x"}
+//   CHECK-DAG:     %[[TID:.+]] = "gpu.thread_id"() {dimension = "x"}
+//       CHECK:     %[[BOFFSET:.+]] = muli %[[BID]], %[[BDIM]]
+//       CHECK:     %[[IV:.+]] = addi %[[BOFFSET]], %[[TID]]
+//       CHECK:     %[[COND:.+]] = cmpi "slt", %[[IV]], %[[C360]]
+//       CHECK:     scf.if %[[COND]]
+//       CHECK:       %[[IV0:.+]] = divi_signed %[[IV]], %[[C120]]
+//       CHECK:       %[[T14:.+]] = remi_signed %[[IV]], %[[C120]]
+//       CHECK:       %[[IV1:.+]] = divi_signed %[[T14]], %[[C30]]
+//       CHECK:       %[[T16:.+]] = remi_signed %[[T14]], %[[C30]]
+//       CHECK:       %[[IV2:.+]] = divi_signed %[[T16]], %[[C6]]
+//       CHECK:       %[[IV3:.+]] = remi_signed %[[T16]], %[[C6]]
+//       CHECK:       load %{{.+}}[%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]]
+//       CHECK:       load %{{.+}}[%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]]
+//       CHECK:       store %{{.+}}[%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]]
+
+//       CHECK: func @[[NUM_WORKGROUPS_FN]]
+//   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//   CHECK-DAG:   %[[C12:.+]] = constant 12 : index
+//       CHECK:   return %[[C12]], %[[C1]], %[[C1]]
+// -----
+
 #map0 = affine_map<() -> ()>
 #accesses = [#map0, #map0, #map0]
 #trait = {
@@ -71,9 +147,13 @@
     [Shader], [SPV_KHR_storage_buffer_storage_class]>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @scalar_add(%arg0 : memref<f32>, %arg1 : memref<f32>,
-                   %arg2 : memref<f32>)
-  {
+  func @scalar_add() attributes {vkspv.num_workgroups_fn = @scalar_add__num_workgroups__} {
+    %arg0 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<f32>
+    %arg1 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg1, operand_result_index = 1 : i32} : memref<f32>
+    %arg2 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@ret0, operand_result_index = 2 : i32} : memref<f32>
     linalg.generic #trait %arg0, %arg1, %arg2 {
     ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
       %0 = addf %arg3, %arg4 : f32
@@ -81,21 +161,38 @@
      } : memref<f32>, memref<f32>, memref<f32>
      return
   }
+  func @scalar_add__num_workgroups__
+    (!shapex.ranked_shape<[]>, !shapex.ranked_shape<[]>,
+     !shapex.ranked_shape<[]>) -> (index, index, index)
+    attributes {sym_visibility = "private"}
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
-// CHECK-LABEL: func @scalar_add
-//  CHECK-SAME:   local_size = dense<1> : vector<3xi32>
-//  CHECK-SAME:   vkspv.workgroup_count_from_result_shape = 1
-//  CHECK-NEXT:     load
+// CHECK-LABEL: func @scalar_add()
+//  CHECK-SAME:   vkspv.num_workgroups_fn = @[[NUM_WORKGROUPS_FN:[a-zA-Z0-9_]+]]
+//       CHECK:     load
 //  CHECK-NEXT:     load
 //  CHECK-NEXT:     addf
 //  CHECK-NEXT:     store
 //  CHECK-NEXT:     return
 
+//       CHECK: func @[[NUM_WORKGROUPS_FN]]
+//   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//       CHECK:   return %[[C1]], %[[C1]], %[[C1]]
+
 // -----
 
 module {
-  func @reduce_sum(%arg0: memref<?x?x?xf32>, %arg1: memref<f32>, %arg2: memref<?xf32>)
-   attributes {iree.dispatch_fn_name = "reduce_sum"} {
+  func @reduce_sum() {
+    %arg0 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<?x?x?xf32>
+    %arg1 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg1, operand_result_index = 1 : i32} : memref<f32>
+    %arg2 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@ret0, operand_result_index = 2 : i32} : memref<?xf32>
     linalg.indexed_generic
       {args_in = 2 : i64, args_out = 1 : i64,
        indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> ()>,
@@ -113,10 +210,14 @@
     }: memref<?x?x?xf32>, memref<f32>, memref<?xf32>
     return
   }
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
 // CHECK-LABEL: func @reduce_sum
 //  CHECK-SAME:   local_size = dense<[32, 1, 1]> : vector<3xi32>
-//  CHECK-SAME:   vkspv.workgroup_count_from_result_shape = 1
 //   CHECK-DAG:     %[[C0:.+]] = constant 0 : index
 //   CHECK-DAG:     %[[C1:.+]] = constant 1 : index
 //   CHECK-DAG:     %[[C2:.+]] = constant 2 : index
@@ -146,10 +247,11 @@
   spv.target_env =
     #spv.target_env<#spv.vce<v1.3, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
                     {max_compute_workgroup_invocations = 128 : i32,
-		     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @matmul(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>)
-  attributes {spv.entry_point_abi = {local_size = dense<[8, 8, 1]> : vector<3xi32>},
-              vkspv.workgroup_count_from_result_shape = 2 : i32} {
+                     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
+  func @matmul() {
+    %arg0 = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg0} : memref<?x?xf32>
+    %arg1 = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg1} : memref<?x?xf32>
+    %arg2 = iree.placeholder for "interace buffer" {binding = @legacy_io::@ret0} : memref<?x?xf32>
     %c4 = constant 4 : index
     %c0 = constant 0 : index
     %c1 = constant 1 : index
@@ -178,6 +280,11 @@
     }
     return
   }
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
 
 // CHECK-LABEL: func @matmul
@@ -209,12 +316,11 @@
   spv.target_env =
     #spv.target_env<#spv.vce<v1.3, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
                     {max_compute_workgroup_invocations = 128 : i32,
-		     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @conv_no_padding(%arg0: memref<?x?x?x?xf32>,
-                        %arg1: memref<?x?x?x?xf32>,
-			%arg2: memref<?x?x?x?xf32>)
-    attributes {spv.entry_point_abi = {local_size = dense<[32, 4, 1]> : vector<3xi32>},
-                vkspv.workgroup_count_from_result_shape = 0 : i32} {
+                     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
+  func @conv_no_padding() {
+    %arg0 = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg0} : memref<?x?x?x?xf32>
+    %arg1 = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg1} : memref<?x?x?x?xf32>
+    %arg2 = iree.placeholder for "interace buffer" {binding = @legacy_io::@ret0} : memref<?x?x?x?xf32>
     %c2 = constant 2 : index
     %c0 = constant 0 : index
     %c3 = constant 3 : index
@@ -252,24 +358,29 @@
               : memref<?x?x?x?xf32> to memref<?x?x?x?xf32, #map5>
       linalg.conv(%arg0, %21, %27)
         {__internal_linalg_transform__ = "workgroup", dilations = [1, 1], strides = [1, 1]}
-	: memref<?x?x?x?xf32>, memref<?x?x?x?xf32, #map5>, memref<?x?x?x?xf32, #map5>
+        : memref<?x?x?x?xf32>, memref<?x?x?x?xf32, #map5>, memref<?x?x?x?xf32, #map5>
       scf.yield
     }
     return
   }
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 * 4)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 * 32)>
 //       CHECK: func @conv_no_padding
-//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9$._-]+]]: memref<?x?x?x?xf32>
-//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9$._-]+]]: memref<?x?x?x?xf32>
-//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9$._-]+]]: memref<?x?x?x?xf32>
+//   CHECK-DAG:   %[[ARG0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg0}
+//   CHECK-DAG:   %[[ARG1:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg1}
+//   CHECK-DAG:   %[[RET0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@ret0}
 //   CHECK-DAG:   %[[C0:.+]] = constant 0
 //   CHECK-DAG:   %[[C1:.+]] = constant 1
 //   CHECK-DAG:   %[[C2:.+]] = constant 2
 //   CHECK-DAG:   %[[N:.+]] = dim %[[ARG1]], %[[C0]]
-//   CHECK-DAG:   %[[P:.+]] = dim %[[ARG2]], %[[C1]]
-//   CHECK-DAG:   %[[Q:.+]] = dim %[[ARG2]], %[[C2]]
+//   CHECK-DAG:   %[[P:.+]] = dim %[[RET0]], %[[C1]]
+//   CHECK-DAG:   %[[Q:.+]] = dim %[[RET0]], %[[C2]]
 //   CHECK-DAG:   %[[BIDX:.+]] = "gpu.block_id"() {dimension = "x"}
 //   CHECK-DAG:   %[[NBLOCKSX:.+]] = "gpu.grid_dim"() {dimension = "x"}
 //   CHECK-DAG:   %[[BIDY:.+]] = "gpu.block_id"() {dimension = "y"}
@@ -284,7 +395,7 @@
 //       CHECK:     scf.for %[[IV4:.+]] = %[[BOFFSETY]] to %[[P]] step %[[BSTEPY]]
 //       CHECK:       scf.for %[[IV5:.+]] = %[[BOFFSETX]] to %[[Q]] step %[[BSTEPX]]
 //       CHECK:         %[[SV1:.+]] = subview %[[ARG1]][%[[IV3]], %[[IV4]], %[[IV5]], 0]
-//       CHECK:         %[[SV2:.+]] = subview %[[ARG2]][%[[IV3]], %[[IV4]], %[[IV5]], 0]
+//       CHECK:         %[[SV2:.+]] = subview %[[RET0]][%[[IV3]], %[[IV4]], %[[IV5]], 0]
 //   CHECK-DAG:         %[[TIDX:.+]] = "gpu.thread_id"() {dimension = "x"}
 //   CHECK-DAG:         %[[NTHREADSX:.+]] = "gpu.block_dim"() {dimension = "x"}
 //   CHECK-DAG:         %[[TIDY:.+]] = "gpu.thread_id"() {dimension = "y"}
@@ -315,10 +426,11 @@
   spv.target_env =
     #spv.target_env<#spv.vce<v1.3, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
                     {max_compute_workgroup_invocations = 128 : i32,
-		     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @pooling_no_padding(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>)
-    attributes {spv.entry_point_abi = {local_size = dense<[32, 4, 1]> : vector<3xi32>},
-                vkspv.workgroup_count_from_result_shape = 0 : i32} {
+                     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
+  func @pooling_no_padding() {
+    %arg0 = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg0} : memref<?x?xf32>
+    %arg1 = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg1} : memref<?x?xf32>
+    %arg2 = iree.placeholder for "interace buffer" {binding = @legacy_io::@ret0} : memref<?x?xf32>
     %c0 = constant 0 : index
     %c1 = constant 1 : index
     %0 = dim %arg1, %c0 : memref<?x?xf32>
@@ -344,23 +456,28 @@
       %19 = subview %arg2[%arg3, %arg4] [%17, %18] [1, 1]  : memref<?x?xf32> to memref<?x?xf32, #map4>
       linalg.pooling_max(%16, %arg1, %19)
         {__internal_linalg_transform__ = "workgroup", dilations = [1, 1], strides = [1, 1]}
-	: memref<?x?xf32, #map4>, memref<?x?xf32>, memref<?x?xf32, #map4>
+        : memref<?x?xf32, #map4>, memref<?x?xf32>, memref<?x?xf32, #map4>
       scf.yield
     }
     return
   }
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
 
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 * 4)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 * 32)>
 //       CHECK: func @pooling_no_padding
-//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9$._-]+]]: memref<?x?xf32>
-//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9$._-]+]]: memref<?x?xf32>
-//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9$._-]+]]: memref<?x?xf32>
+//   CHECK-DAG:   %[[ARG0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg0}
+//   CHECK-DAG:   %[[ARG1:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg1}
+//   CHECK-DAG:   %[[RET0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@ret0}
 //   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
 //   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
-//   CHECK-DAG:   %[[P:.+]] = dim %[[ARG2]], %[[C0]]
-//   CHECK-DAG:   %[[Q:.+]] = dim %[[ARG2]], %[[C1]]
+//   CHECK-DAG:   %[[P:.+]] = dim %[[RET0]], %[[C0]]
+//   CHECK-DAG:   %[[Q:.+]] = dim %[[RET0]], %[[C1]]
 //   CHECK-DAG:   %[[BIDX:.+]] = "gpu.block_id"() {dimension = "x"}
 //   CHECK-DAG:   %[[NBLOCKSX:.+]] = "gpu.grid_dim"() {dimension = "x"}
 //   CHECK-DAG:   %[[BIDY:.+]] = "gpu.block_id"() {dimension = "y"}
@@ -372,7 +489,7 @@
 //       CHECK:   scf.for %[[IV3:.+]] = %[[BOFFSETY]] to %[[P]] step %[[BSTEPY]]
 //       CHECK:     scf.for %[[IV4:.+]] = %[[BOFFSETX]] to %[[Q]] step %[[BSTEPX]]
 //       CHECK:       %[[SV1:.+]] = subview %[[ARG0]][%[[IV3]], %[[IV4]]]
-//       CHECK:       %[[SV2:.+]] = subview %[[ARG2]][%[[IV3]], %[[IV4]]]
+//       CHECK:       %[[SV2:.+]] = subview %[[RET0]][%[[IV3]], %[[IV4]]]
 //   CHECK-DAG:       %[[TIDX:.+]] = "gpu.thread_id"() {dimension = "x"}
 //   CHECK-DAG:       %[[NTHREADSX:.+]] = "gpu.block_dim"() {dimension = "x"}
 //   CHECK-DAG:       %[[TIDY:.+]] = "gpu.thread_id"() {dimension = "y"}
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/linalg_tile_and_fuse.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/linalg_tile_and_fuse.mlir
index 33cfde8..e06eae5 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/test/linalg_tile_and_fuse.mlir
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/linalg_tile_and_fuse.mlir
@@ -7,23 +7,30 @@
     [Shader], [SPV_KHR_storage_buffer_storage_class]>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @conv_padding(%arg0 : memref<?x?x?x?xf32>, %arg1 : memref<?x?x?x?xf32>,
-                     %arg2 : memref<?x?x?x?xf32>) {
-    linalg.conv(%arg0, %arg1, %arg2)
+  func @conv_padding() {
+    %0 = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg0} : memref<?x?x?x?xf32>
+    %1 = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg1} : memref<?x?x?x?xf32>
+    %2 = iree.placeholder for "interace buffer" {binding = @legacy_io::@ret0} : memref<?x?x?x?xf32>
+    linalg.conv(%0, %1, %2)
       {dilations = [1, 1],
        padding = dense<[[1, 1], [0, 1]]> : tensor<2x2xi64>, strides = [1, 1]} :
       memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
     return
   }
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
-// CHECK-LABEL: func @conv_padding
-//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9$._-]+]]: memref<?x?x?x?xf32>
-//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9$._-]+]]: memref<?x?x?x?xf32>
-//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9$._-]+]]: memref<?x?x?x?xf32>
+//       CHECK: func @conv_padding()
+//   CHECK-DAG:   %[[ARG0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg0}
+//   CHECK-DAG:   %[[ARG1:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg1}
+//   CHECK-DAG:   %[[RET0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@ret0}
 //       CHECK:   linalg.conv
 //  CHECK-SAME:     %[[ARG0]]
 //  CHECK-SAME:     %[[ARG1]]
-//  CHECK-SAME:     %[[ARG2]]
+//  CHECK-SAME:     %[[RET0]]
 
 // -----
 
@@ -33,21 +40,30 @@
     [Shader], [SPV_KHR_storage_buffer_storage_class]>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @conv_no_padding(%arg0 : memref<?x?x?x?xf32>, %arg1 : memref<?x?x?x?xf32>,
-                        %arg2 : memref<?x?x?x?xf32>) {
-    linalg.conv(%arg0, %arg1, %arg2) {dilations = [1, 1], strides = [1, 1]} :
+  func @conv_no_padding() {
+    %0 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<?x?x?x?xf32>
+    %1 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg1, operand_result_index = 1 : i32} : memref<?x?x?x?xf32>
+    %2 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@ret0, operand_result_index = 2 : i32} : memref<?x?x?x?xf32>
+    linalg.conv(%0, %1, %2) {dilations = [1, 1], strides = [1, 1]} :
       memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
     return
   }
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 * 4)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 * 32)>
-//       CHECK: func @conv_no_padding
-//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9$._-]+]]: memref<?x?x?x?xf32>
-//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9$._-]+]]: memref<?x?x?x?xf32>
-//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9$._-]+]]: memref<?x?x?x?xf32>
+//       CHECK: func @conv_no_padding()
 //  CHECK-SAME:   local_size = dense<[32, 4, 1]>
-//  CHECK-SAME:   vkspv.workgroup_count_from_result_shape = 0
+//   CHECK-DAG:   %[[ARG0:.+]] = iree.placeholder {{.*}} {binding = @legacy_io::@arg0
+//   CHECK-DAG:   %[[ARG1:.+]] = iree.placeholder {{.*}} {binding = @legacy_io::@arg1
+//   CHECK-DAG:   %[[RET0:.+]] = iree.placeholder {{.*}} {binding = @legacy_io::@ret0
 //   CHECK-DAG:   %[[C0:.+]] = constant 0
 //   CHECK-DAG:   %[[BIDX:.+]] = "gpu.block_id"() {dimension = "x"}
 //   CHECK-DAG:   %[[NBLOCKSX:.+]] = "gpu.grid_dim"() {dimension = "x"}
@@ -59,12 +75,13 @@
 //       CHECK:   %[[STEPY:.+]] = affine.apply #[[MAP0]]()[%[[NBLOCKSY]]]
 //       CHECK:   %[[LBX:.+]] = affine.apply #[[MAP1]]()[%[[BIDX]]
 //       CHECK:   %[[STEPX:.+]] = affine.apply #[[MAP1]]()[%[[NBLOCKSX]]]
-//       CHECK:   scf.parallel (%[[IV0:.+]], %[[IV1:.+]], %[[IV2:.+]]) =
-//       (%[[BIDZ]], %[[LBY]], %[[LBX]])
-//  CHECK-SAME:      step (%[[NBLOCKSZ]], %[[STEPY]], %[[STEPX]])
-//       CHECK:     %[[VIEW1:.+]] = subview %[[ARG1]][%[[IV0]], %[[IV1]],
-//       %[[IV2]], %[[C0]]] CHECK:     %[[VIEW2:.+]] = subview
-//       %[[ARG2]][%[[IV0]], %[[IV1]], %[[IV2]], %[[C0]]] CHECK:     linalg.conv
+//       CHECK:   scf.parallel (%[[IV0:.+]], %[[IV1:.+]], %[[IV2:.+]]) = (%[[BIDZ]], %[[LBY]], %[[LBX]])
+//  CHECK-SAME:     step (%[[NBLOCKSZ]], %[[STEPY]], %[[STEPX]])
+//       CHECK:     %[[VIEW1:.+]] = subview %[[ARG1]]
+//  CHECK-SAME:       [%[[IV0]], %[[IV1]], %[[IV2]], %[[C0]]]
+//       CHECK:     %[[VIEW2:.+]] = subview %[[RET0]]
+//  CHECK-SAME:       [%[[IV0]], %[[IV1]], %[[IV2]], %[[C0]]]
+//       CHECK:     linalg.conv
 //  CHECK-SAME:       %[[ARG0]], %[[VIEW1]], %[[VIEW2]]
 //  CHECK-SAME:       "workgroup"
 
@@ -76,22 +93,35 @@
     [Shader], [SPV_KHR_storage_buffer_storage_class]>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @matmul(%arg0: memref<?x?xf32>,
-                %arg1: memref<?x?xf32>,
-                %ret0: memref<?x?xf32>) {
-    linalg.matmul %arg0, %arg1, %ret0 :
+  func @matmul() attributes {vkspv.num_workgroups_fn = @matmul__num_workgroups__} {
+    %0 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<?x?xf32>
+    %1 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg1, operand_result_index = 1 : i32} : memref<?x?xf32>
+    %2 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@ret0, operand_result_index = 2 : i32} : memref<?x?xf32>
+    linalg.matmul %0, %1, %2 :
       (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
     return
   }
+  func @matmul__num_workgroups__
+    (!shapex.ranked_shape<[?,?]>, !shapex.ranked_shape<[?,?]>,
+     !shapex.ranked_shape<[?,?]>) -> (index, index, index)
+    attributes {sym_visibility = "private"}
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
 
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 * 8)>
-//       CHECK: func @matmul
-//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9$._-]+]]: memref<?x?xf32>
-//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9$._-]+]]: memref<?x?xf32>
-//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9$._-]+]]: memref<?x?xf32>
+//       CHECK: func @matmul()
 //  CHECK-SAME:   local_size = dense<[8, 8, 1]>
-//  CHECK-SAME:   vkspv.workgroup_count_from_result_shape = 2
+//  CHECK-SAME:   vkspv.num_workgroups_fn = @[[NUM_WORKGROUPS_FN:.[a-zA-Z0-9_]+]]
+//   CHECK-DAG:   %[[ARG0:.+]] = iree.placeholder {{.*}} {binding = @legacy_io::@arg0
+//   CHECK-DAG:   %[[ARG1:.+]] = iree.placeholder {{.*}} {binding = @legacy_io::@arg1
+//   CHECK-DAG:   %[[RET0:.+]] = iree.placeholder {{.*}} {binding = @legacy_io::@ret0
 //   CHECK-DAG:   %[[C0:.+]] = constant 0
 //   CHECK-DAG:   %[[C4:.+]] = constant 4
 //   CHECK-DAG:   %[[BIDX:.+]] = "gpu.block_id"() {dimension = "x"}
@@ -104,10 +134,22 @@
 //       CHECK:     %[[VIEW1:.+]] = subview %[[ARG1]][%[[IV]], %[[LBX]]]
 //       CHECK:     %[[LBY_2:.+]] = affine.apply #[[MAP0]]()[%[[BIDY]]]
 //       CHECK:     %[[LBX_2:.+]] = affine.apply #[[MAP0]]()[%[[BIDX]]]
-//       CHECK:     %[[VIEW2:.+]] = subview %[[ARG2]][%[[LBY_2]], %[[LBX_2]]]
+//       CHECK:     %[[VIEW2:.+]] = subview %[[RET0]][%[[LBY_2]], %[[LBX_2]]]
 //       CHECK:     linalg.matmul
 //  CHECK-SAME:       "workgroup_numprocs_ge_numiters"
 //  CHECK-SAME:       %[[VIEW0]], %[[VIEW1]], %[[VIEW2]]
+//       CHECK: func @[[NUM_WORKGROUPS_FN]]
+//   CHECK-DAG:   %[[C8:.+]] = constant 8 : index
+//   CHECK-DAG:   %[[C7:.+]] = constant 7 : index
+//   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
+//   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//       CHECK:   %[[DIM0:.+]] = dim %{{.*}}, %[[C0]]
+//       CHECK:   %[[DIM1:.+]] = dim %{{.*}}, %[[C1]]
+//       CHECK:   %[[T0:.+]] = addi %[[DIM1]], %[[C7]]
+//       CHECK:   %[[T1:.+]] = divi_signed %[[T0]], %[[C8]]
+//       CHECK:   %[[T2:.+]] = addi %[[DIM0]], %[[C7]]
+//       CHECK:   %[[T3:.+]] = divi_signed %[[T2]], %[[C8]]
+//       CHECK:   return %[[T1]], %[[T3]], %[[C1]]
 
 // -----
 
@@ -117,21 +159,30 @@
     [Shader], [SPV_KHR_storage_buffer_storage_class]>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @pooling_sum_no_padding(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>,
-                               %arg2 : memref<?x?xf32>) {
-    linalg.pooling_max(%arg0, %arg1, %arg2) {dilations = [1, 1], strides = [1, 1]} :
+  func @pooling_sum_no_padding() {
+    %0 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<?x?xf32>
+    %1 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg1, operand_result_index = 1 : i32} : memref<?x?xf32>
+    %2 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@ret0, operand_result_index = 2 : i32} : memref<?x?xf32>
+    linalg.pooling_max(%0, %1, %2) {dilations = [1, 1], strides = [1, 1]} :
       memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>
     return
   }
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 * 4)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 * 32)>
-//       CHECK: func @pooling_sum_no_padding
-//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9$._-]+]]: memref<?x?xf32>
-//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9$._-]+]]: memref<?x?xf32>
-//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9$._-]+]]: memref<?x?xf32>
+//       CHECK: func @pooling_sum_no_padding()
 //  CHECK-SAME:   local_size = dense<[32, 4, 1]>
-//  CHECK-SAME:   vkspv.workgroup_count_from_result_shape = 0
+//   CHECK-DAG:   %[[ARG0:.+]] = iree.placeholder {{.*}} {binding = @legacy_io::@arg0
+//   CHECK-DAG:   %[[ARG1:.+]] = iree.placeholder {{.*}} {binding = @legacy_io::@arg1
+//   CHECK-DAG:   %[[RET0:.+]] = iree.placeholder {{.*}} {binding = @legacy_io::@ret0
 //   CHECK-DAG:   %[[BIDX:.+]] = "gpu.block_id"() {dimension = "x"}
 //   CHECK-DAG:   %[[NBLOCKSX:.+]] = "gpu.grid_dim"() {dimension = "x"}
 //   CHECK-DAG:   %[[BIDY:.+]] = "gpu.block_id"() {dimension = "y"}
@@ -143,7 +194,7 @@
 //       CHECK:   scf.parallel (%[[IV0:.+]], %[[IV1:.+]]) = (%[[LBY]], %[[LBX]])
 //  CHECK-SAME:      step (%[[STEPY]], %[[STEPX]])
 //       CHECK:     %[[VIEW0:.+]] = subview %[[ARG0]][%[[IV0]], %[[IV1]]]
-//       CHECK:     %[[VIEW2:.+]] = subview %[[ARG2]][%[[IV0]], %[[IV1]]]
+//       CHECK:     %[[VIEW2:.+]] = subview %[[RET0]][%[[IV0]], %[[IV1]]]
 //       CHECK:     linalg.pooling_max
 //  CHECK-SAME:       %[[VIEW0]], %[[ARG1]], %[[VIEW2]]
 //  CHECK-SAME:       "workgroup"
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/matmul_vectorization.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/matmul_vectorization.mlir
index 63ba65f..c642f62 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/test/matmul_vectorization.mlir
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/matmul_vectorization.mlir
@@ -9,13 +9,13 @@
     return
 }
 
-// CHECK: %[[TILESIZE:.+]] = constant 32 : index
-// CHECK: %[[MATSIZE:.+]] = constant 128 : index
-// CHECK: %[[START:.+]] = constant 0 : index
-// CHECK: scf.for %[[IL:.+]] = %[[START]] to %[[MATSIZE]] step %[[TILESIZE]]
-// CHECK: scf.for %[[JL:.+]] = %[[START]] to %[[MATSIZE]] step %[[TILESIZE]]
-// CHECK: %[[SUBVVIEWC:.+]] = subview %[[ARG2]][%[[IL]], %[[JL]]] [32, 32] [1, 1] : memref<128x128xf32> to memref<32x32xf32
-// CHECK: scf.for %[[KL:.+]] = %[[START]] to %[[MATSIZE]] step %[[TILESIZE]]
-// CHECK: %[[SUBVVIEWA:.+]] = subview %[[ARG0]][%[[IL]], %[[KL]]] [32, 32] [1, 1] : memref<128x128xf32> to memref<32x32xf32
-// CHECK: %[[SUBVVIEWB:.+]] = subview %[[ARG1]][%[[KL]], %[[JL]]] [32, 32] [1, 1] : memref<128x128xf32> to memref<32x32xf32
+// CHECK-DAG: %[[TILESIZE:.+]] = constant 32 : index
+// CHECK-DAG: %[[MATSIZE:.+]] = constant 128 : index
+// CHECK-DAG: %[[START:.+]] = constant 0 : index
+//     CHECK: scf.for %[[IL:.+]] = %[[START]] to %[[MATSIZE]] step %[[TILESIZE]]
+//     CHECK:   scf.for %[[JL:.+]] = %[[START]] to %[[MATSIZE]] step %[[TILESIZE]]
+//     CHECK:     %[[SUBVVIEWC:.+]] = subview %[[ARG2]][%[[IL]], %[[JL]]] [32, 32] [1, 1] : memref<128x128xf32> to memref<32x32xf32
+//     CHECK:     scf.for %[[KL:.+]] = %[[START]] to %[[MATSIZE]] step %[[TILESIZE]]
+//     CHECK:       %[[SUBVVIEWA:.+]] = subview %[[ARG0]][%[[IL]], %[[KL]]] [32, 32] [1, 1] : memref<128x128xf32> to memref<32x32xf32
+//     CHECK:       %[[SUBVVIEWB:.+]] = subview %[[ARG1]][%[[KL]], %[[JL]]] [32, 32] [1, 1] : memref<128x128xf32> to memref<32x32xf32
 
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/split_dispatch_function.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/split_dispatch_function.mlir
index dbffa87..c115b4a 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/test/split_dispatch_function.mlir
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/split_dispatch_function.mlir
@@ -23,7 +23,7 @@
   // CHECK:   linalg.fill(%[[TS]], %[[ZERO]])
   // CHECK:   return
 
-  func @kernel() {
+  func @kernel() attributes {vkspv.num_workgroups_fn = @kernel__num_workgroups__} {
     %cst = constant 0.000000e+00 : f32
     %dim = hal.interface.load.constant offset = 0 : index
     %shape1 = shapex.make_ranked_shape %dim : (index) -> !shapex.ranked_shape<[?,2,2,512]>
@@ -37,6 +37,11 @@
     linalg.conv(%1, %ts1, %ts2) {dilations = [1, 1], padding = dense<[[0, 1], [0, 1]]> : tensor<2x2xi64>, strides = [2, 2]} : memref<3x3x512x1xf32>, memref<?x2x2x512xf32>, memref<?x1x1x512xf32>
     return
   }
+  func @kernel__num_workgroups__(!shapex.ranked_shape<[?,2,2,512]>,
+                                 !shapex.ranked_shape<[3,3,512,1]>,
+                                 !shapex.ranked_shape<[?,1,1,512]>)
+                                 -> (index, index, index)
+  attributes {sym_visibility = "private"}
   hal.interface @legacy_io attributes {push_constants = 1 : i32, sym_visibility = "private"} {
     hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
     hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
@@ -48,35 +53,44 @@
 
 // CHECK: module attributes {vkspv.entry_point_schedule = ["kernel_dispatch_0", "kernel_dispatch_1", "kernel_dispatch_2"]}
 module {
-  // CHECK: func @kernel_dispatch_2()
-  // CHECK:   %[[DIM:.+]] = hal.interface.load.constant
-  // CHECK:   %[[SHAPE1:.+]] = shapex.make_ranked_shape %[[DIM]]
-  // CHECK:   %[[SHAPE2:.+]] = shapex.make_ranked_shape %[[DIM]]
-  // CHECK:   %[[IN1:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<?x2x2x512xf32>
-  // CHECK:   %[[TS1:.+]] = shapex.tie_shape %[[IN1]], %[[SHAPE1]]
-  // CHECK:   %[[IN2:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg1} : memref<3x3x512x1xf32>
-  // CHECK:   %[[OUT:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<?x1x1x512xf32>
-  // CHECK:   %[[TS2:.+]] = shapex.tie_shape %[[OUT]], %[[SHAPE2]]
-  // CHECK:   linalg.conv(%[[IN2]], %[[TS1]], %[[TS2]])
-  // CHECK:   return
+//      CHECK: func @kernel_dispatch_2()
+// CHECK-SAME: {vkspv.num_workgroups_fn = @[[NUM_WORKGROUPS_FN2:.+]]}
+//      CHECK:   %[[DIM:.+]] = hal.interface.load.constant
+//      CHECK:   %[[SHAPE1:.+]] = shapex.make_ranked_shape %[[DIM]]
+//      CHECK:   %[[SHAPE2:.+]] = shapex.make_ranked_shape %[[DIM]]
+//      CHECK:   %[[IN1:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<?x2x2x512xf32>
+//      CHECK:   %[[TS1:.+]] = shapex.tie_shape %[[IN1]], %[[SHAPE1]]
+//      CHECK:   %[[IN2:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg1} : memref<3x3x512x1xf32>
+//      CHECK:   %[[OUT:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<?x1x1x512xf32>
+//      CHECK:   %[[TS2:.+]] = shapex.tie_shape %[[OUT]], %[[SHAPE2]]
+//      CHECK:   linalg.conv(%[[IN2]], %[[TS1]], %[[TS2]])
+//      CHECK:   return
 
-  // CHECK: func @kernel_dispatch_1() {
-  // CHECK:   %[[C0:.+]] = constant 0 : index
-  // CHECK:   %[[C1:.+]] = constant 1 : index
-  // CHECK:   scf.parallel (%{{.*}}) = (%[[C0]]) to (%[[C1]]) step (%[[C1]])
-  // CHECK:     scf.yield
-  // CHECK:   return
+//      CHECK: func @[[NUM_WORKGROUPS_FN2]]
 
-  // CHECK: func @kernel_dispatch_0()
-  // CHECK:   %[[ZERO:.+]] = constant
-  // CHECK:   %[[DIM:.+]] = hal.interface.load.constant
-  // CHECK:   %[[SHAPE:.+]] = shapex.make_ranked_shape %[[DIM]]
-  // CHECK:   %[[OUT:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<?x1x1x512xf32>
-  // CHECK:   %[[TS:.+]] = shapex.tie_shape %[[OUT]], %[[SHAPE]]
-  // CHECK:   linalg.fill(%[[TS]], %[[ZERO]])
-  // CHECK:   return
+//      CHECK: func @kernel_dispatch_1()
+// CHECK-SAME: {vkspv.num_workgroups_fn = @[[NUM_WORKGROUPS_FN1:.+]]}
+//      CHECK:   %[[C0:.+]] = constant 0 : index
+//      CHECK:   %[[C1:.+]] = constant 1 : index
+//      CHECK:   scf.parallel (%{{.*}}) = (%[[C0]]) to (%[[C1]]) step (%[[C1]])
+//      CHECK:     scf.yield
+//      CHECK:   return
 
-  func @kernel() {
+//      CHECK: func @[[NUM_WORKGROUPS_FN1]]
+
+//      CHECK: func @kernel_dispatch_0()
+// CHECK-SAME: {vkspv.num_workgroups_fn = @[[NUM_WORKGROUPS_FN0:.+]]}
+//      CHECK:   %[[ZERO:.+]] = constant
+//      CHECK:   %[[DIM:.+]] = hal.interface.load.constant
+//      CHECK:   %[[SHAPE:.+]] = shapex.make_ranked_shape %[[DIM]]
+//      CHECK:   %[[OUT:.+]] = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<?x1x1x512xf32>
+//      CHECK:   %[[TS:.+]] = shapex.tie_shape %[[OUT]], %[[SHAPE]]
+//      CHECK:   linalg.fill(%[[TS]], %[[ZERO]])
+//      CHECK:   return
+
+//      CHECK: func @[[NUM_WORKGROUPS_FN0]]
+
+  func @kernel() attributes {vkspv.num_workgroups_fn = @kernel__num_workgroups__} {
     %cst = constant 0.000000e+00 : f32
     %c0 = constant 0 : index
     %c1 = constant 1 : index
@@ -95,6 +109,11 @@
     linalg.conv(%1, %ts1, %ts2) {dilations = [1, 1], padding = dense<[[0, 1], [0, 1]]> : tensor<2x2xi64>, strides = [2, 2]} : memref<3x3x512x1xf32>, memref<?x2x2x512xf32>, memref<?x1x1x512xf32>
     return
   }
+  func @kernel__num_workgroups__(!shapex.ranked_shape<[?,2,2,512]>,
+                                 !shapex.ranked_shape<[3,3,512,1]>,
+                                 !shapex.ranked_shape<[?,1,1,512]>)
+                                 -> (index, index, index)
+  attributes {sym_visibility = "private"}
   hal.interface @legacy_io attributes {push_constants = 1 : i32, sym_visibility = "private"} {
     hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
     hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
@@ -110,7 +129,7 @@
 // CHECK-NOT: vkspv.entry_point_schedule
 module {
   // CHECK-LABEL: @kernel()
-  func @kernel() {
+  func @kernel() attributes {vkspv.num_workgroups_fn = @kernel__num_workgroups__} {
     %cst = constant 0.000000e+00 : f32
     %0 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<1x2x2x512xf32>
     %1 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg1} : memref<3x3x512x1xf32>
@@ -118,6 +137,8 @@
     linalg.conv(%1, %0, %2) {dilations = [1, 1], padding = dense<[[0, 1], [0, 1]]> : tensor<2x2xi64>, strides = [2, 2]} : memref<3x3x512x1xf32>, memref<1x2x2x512xf32>, memref<1x1x1x512xf32>
     return
   }
+  // CHECK-LABEL: @kernel__num_workgroups__
+  func @kernel__num_workgroups__(!shapex.ranked_shape<[1,2,2,512]>, !shapex.ranked_shape<[3,3,1,512]>, !shapex.ranked_shape<[1,1,1,512]>) -> (index, index, index) attributes {sym_visibility = "private"}
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
     hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
     hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
@@ -153,7 +174,8 @@
 #map0 = affine_map<(d0, d1) -> (d0 * 12 + d1 + 53)>
 
 module {
-  func @subview_interleaved() {
+  func @subview_interleaved()
+  attributes {vkspv.num_workgroups_fn = @subview_interleaved__num_workgroups__} {
     %cst = constant 0.000000e+00 : f32
     %0 = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<18x12xf32>
     %1 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<12x4xf32>
@@ -162,6 +184,10 @@
     linalg.copy(%1, %2) : memref<12x4xf32>, memref<18x12xf32, #map0>
     return
   }
+  func @subview_interleaved__num_workgroups__(!shapex.ranked_shape<[12,4]>,
+                                              !shapex.ranked_shape<[18,12]>)
+                                              -> (index, index, index)
+  attributes {sym_visibility = "private"}
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
     hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
     hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer", access="Write"
@@ -191,7 +217,8 @@
 #map2 = affine_map<(d0, d1, d2) -> (d2)>
 
 module {
-  func @reshape_interleaved() {
+  func @reshape_interleaved()
+  attributes {vkspv.num_workgroups_fn = @reshape_interleaved__num_workgroups__} {
     %0 = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<2x4xf32>
     %1 = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret1} : memref<1x2x4xf32>
     %2 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<2x4xf32>
@@ -206,6 +233,11 @@
     linalg.copy(%3, %1) : memref<1x2x4xf32>, memref<1x2x4xf32>
     return
   }
+  func @reshape_interleaved__num_workgroups__(!shapex.ranked_shape<[2,4]>,
+                                              !shapex.ranked_shape<[2,4]>,
+                                              !shapex.ranked_shape<[1,2,4]>)
+                                              -> (index, index, index)
+  attributes {sym_visibility = "private"}
   hal.interface @legacy_io attributes {sym_visibility = "private"} {
     hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
     hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer", access="Write|Discard"
@@ -233,7 +265,8 @@
 // -----
 
 module {
-  func @predict_ex_dispatch_0() {
+  func @predict_ex_dispatch_0()
+  attributes {vkspv.num_workgroups_fn = @predict_ex_dispatch_0__num_workgroups__} {
     %0 = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret0} : memref<1x512x1xf32>
     %1 = iree.placeholder for "interface buffer" {binding = @legacy_io::@ret1} : memref<4x8x16xf32>
     %2 = iree.placeholder for "interface buffer" {binding = @legacy_io::@arg0} : memref<1x512x1xf32>
@@ -248,6 +281,12 @@
     }: memref<4x8x16xf32>, memref<4x8x16xf32>
     return
   }
+  func @predict_ex_dispatch_0__num_workgroups__(!shapex.ranked_shape<[1,512,1]>,
+                                                !shapex.ranked_shape<[4,8,16]>,
+                                                !shapex.ranked_shape<[1,512,1]>,
+                                                !shapex.ranked_shape<[4,8,16]>)
+                                                -> (index, index, index)
+  attributes {sym_visibility = "private"}
   hal.interface @legacy_io attributes {push_constants = 1 : i32, sym_visibility = "private"} {
     hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
     hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/test/workgroup_memory_promotion.mlir b/iree/compiler/Conversion/LinalgToSPIRV/test/workgroup_memory_promotion.mlir
index fbbb483..9ff21ba 100644
--- a/iree/compiler/Conversion/LinalgToSPIRV/test/workgroup_memory_promotion.mlir
+++ b/iree/compiler/Conversion/LinalgToSPIRV/test/workgroup_memory_promotion.mlir
@@ -4,21 +4,32 @@
   spv.target_env =
     #spv.target_env<#spv.vce<v1.3, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
                     {max_compute_workgroup_invocations = 128 : i32,
-         max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @matmul_tile(%arg0 : memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
-    linalg.matmul %arg0, %arg1, %arg2 :
-      (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
+                     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
+  func @matmul_tile()
+    attributes {signature = (tensor<?x?xf32>, tensor<?x?xf32>) -> (tensor<?x?xf32>)} {
+    %0 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<?x?xf32>
+    %1 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg1, operand_result_index = 1 : i32} : memref<?x?xf32>
+    %2 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@ret0, operand_result_index = 2 : i32} : memref<?x?xf32>
+    linalg.matmul %0, %1, %2 : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
     return
   }
+  hal.interface @legacy_io attributes {sym_visibility = "private"} {
+    hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read"
+    hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read"
+    hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write"
+  }
 }
-// CHECK-LABEL: func @matmul_tile
-//  CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?xf32>
-//  CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?xf32>
-//  CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+//       CHECK: func @matmul_tile()
+//   CHECK-DAG:   %[[ARG0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg0
+//   CHECK-DAG:   %[[ARG1:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg1
+//   CHECK-DAG:   %[[RET0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@ret0
 //       CHECK:   scf.for
 //       CHECK:     %[[ARG0SV:.+]] = subview %[[ARG0]]
 //       CHECK:     %[[ARG1SV:.+]] = subview %[[ARG1]]
-//       CHECK:     %[[ARG2SV:.+]] = subview %[[ARG2]]
+//       CHECK:     %[[RET0SV:.+]] = subview %[[RET0]]
 //       CHECK:     %[[ALLOC1:.+]] = alloc() : memref<8x4xf32, 3>
 //       CHECK:     %[[SUBVIEW1:.+]] = subview %[[ALLOC1]]
 //       CHECK:     %[[ALLOC2:.+]] = alloc() : memref<4x8xf32, 3>
@@ -29,7 +40,7 @@
 //  CHECK-SAME:       "copy_to_workgroup_memory"
 //       CHECK:     linalg.matmul
 //  CHECK-SAME:       "workgroup_memory_numprocs_ge_numiters"
-//  CHECK-SAME:       %[[SUBVIEW1]], %[[SUBVIEW2]], %[[ARG2SV]]
+//  CHECK-SAME:       %[[SUBVIEW1]], %[[SUBVIEW2]], %[[RET0SV]]
 //   CHECK-DAG:     dealloc %[[ALLOC1]] : memref<8x4xf32, 3>
 //   CHECK-DAG:     dealloc %[[ALLOC2]] : memref<4x8xf32, 3>
 
@@ -39,25 +50,31 @@
   spv.target_env =
     #spv.target_env<#spv.vce<v1.3, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
                     {max_compute_workgroup_invocations = 128 : i32,
-         max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
-  func @conv_no_padding_tile(%arg0: memref<3x4x3x2xf32>,
-                             %arg1: memref<?x?x?x3xf32>, %arg2: memref<?x?x?x2xf32>) {
-    linalg.conv(%arg0, %arg1, %arg2) {dilations = [1, 1], strides = [1, 1]}
+                     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>} {
+  func @conv_no_padding_tile()
+    attributes {signature = (tensor<3x4x3x2xf32>, tensor<?x?x?x3xf32>) -> (tensor<?x?x?x2xf32>)}  {
+    %0 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg0, operand_result_index = 0 : i32} : memref<3x4x3x2xf32>
+    %1 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@arg1, operand_result_index = 1 : i32} : memref<?x?x?x3xf32>
+    %2 = iree.placeholder for "interace buffer"
+      {binding = @legacy_io::@ret0, operand_result_index = 2 : i32} : memref<?x?x?x2xf32>
+    linalg.conv(%0, %1, %2) {dilations = [1, 1], strides = [1, 1]}
       : memref<3x4x3x2xf32>, memref<?x?x?x3xf32>, memref<?x?x?x2xf32>
     return
   }
 }
-// CHECK-LABEL: func @conv_no_padding_tile
-//  CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<3x4x3x2xf32>
-//  CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?x?x3xf32>
-//  CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?x?x2xf32>
-//       CHECK: scf.parallel (%{{.*}}, %{{.*}}, %{{.*}})
-//       CHECK:    %[[ARG1SV:.+]] = subview %[[ARG1]]
-//       CHECK:    %[[ARG2SV:.+]] = subview %[[ARG2]]
-//       CHECK:    %[[ALLOC1:.+]] = alloc()
-//       CHECK:    %[[SUBVIEW1:.+]] = subview %[[ALLOC1]]
-//       CHECK:    linalg.copy(%[[ARG1SV]], %[[SUBVIEW1]])
-//  CHECK-SAME:       "copy_to_workgroup_memory"
-//       CHECK:    linalg.conv(%[[ARG0]], %[[SUBVIEW1]], %[[ARG2SV]])
-//  CHECK-SAME:       "workgroup_memory"
-//       CHECK:    dealloc %[[ALLOC1]]
+//       CHECK: func @conv_no_padding_tile()
+//   CHECK-DAG:   %[[ARG0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg0
+//   CHECK-DAG:   %[[ARG1:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@arg1
+//   CHECK-DAG:   %[[RET0:.+]] = iree.placeholder for "interace buffer" {binding = @legacy_io::@ret0
+//       CHECK:   scf.parallel (%{{.*}}, %{{.*}}, %{{.*}})
+//       CHECK:     %[[ARG1SV:.+]] = subview %[[ARG1]]
+//       CHECK:     %[[RET0SV:.+]] = subview %[[RET0]]
+//       CHECK:     %[[ALLOC1:.+]] = alloc() : memref<1x6x35x3xf32, 3>
+//       CHECK:     %[[SUBVIEW1:.+]] = subview %[[ALLOC1]]
+//       CHECK:     linalg.copy(%[[ARG1SV]], %[[SUBVIEW1]])
+//  CHECK-SAME:        "copy_to_workgroup_memory"
+//       CHECK:     linalg.conv(%[[ARG0]], %[[SUBVIEW1]], %[[RET0SV]])
+//  CHECK-SAME:        "workgroup_memory"
+//       CHECK:     dealloc %[[ALLOC1]] : memref<1x6x35x3xf32, 3>
diff --git a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/BUILD b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/BUILD
index 18006a7..97f0662 100644
--- a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/BUILD
+++ b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/BUILD
@@ -41,6 +41,7 @@
         "//iree/compiler/Conversion/LinalgToSPIRV",
         "//iree/compiler/Dialect/Flow/IR",
         "//iree/compiler/Dialect/HAL/Target",
+        "//iree/compiler/Dialect/Shape/IR",
         "//iree/compiler/Dialect/Vulkan/IR",
         "//iree/compiler/Dialect/Vulkan/Utils",
         "//iree/schemas:spirv_executable_def_cc_fbs",
diff --git a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/CMakeLists.txt
index 05d2710..fdc9911 100644
--- a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/CMakeLists.txt
@@ -44,6 +44,7 @@
     iree::compiler::Conversion::LinalgToSPIRV
     iree::compiler::Dialect::Flow::IR
     iree::compiler::Dialect::HAL::Target
+    iree::compiler::Dialect::Shape::IR
     iree::compiler::Dialect::Vulkan::IR
     iree::compiler::Dialect::Vulkan::Utils
     iree::schemas::spirv_executable_def_cc_fbs
diff --git a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
index 6851ec2..f18c2b8 100644
--- a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
@@ -21,6 +21,7 @@
 #include "iree/compiler/Conversion/LinalgToSPIRV/Attributes.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
 #include "iree/compiler/Dialect/Vulkan/IR/VulkanAttributes.h"
 #include "iree/compiler/Dialect/Vulkan/IR/VulkanDialect.h"
 #include "iree/compiler/Dialect/Vulkan/Utils/TargetEnvUtils.h"
@@ -37,6 +38,7 @@
 #include "mlir/Dialect/SPIRV/Serialization.h"
 #include "mlir/Dialect/SPIRV/TargetAndABI.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Matchers.h"
@@ -127,105 +129,61 @@
       ArrayRef<Value>{memoryBarrier}, ArrayRef<Value>{});
 }
 
-/// Generates IR to compute the ceil(`numerator`, `denominator`).
-static Value computeCeilDiv(Location loc, Value one, Value numerator,
-                            Value denominator, OpBuilder &builder) {
-  Value dm1 = builder.create<SubIOp>(loc, denominator, one);
-  return builder.create<SignedDivIOp>(
-      loc, builder.create<AddIOp>(loc, numerator, dm1), denominator);
-}
-
-/// Calculates the number of workgroups to use based on the shape of the result
-/// of the dispatch region.  If the `resultShape` is {s0, s1, s2, s3, ....} and
-/// `workgroupSize` is {wx, wy, wz}, the number of workgroups is {ceil(s0/wz),
-/// ceil(s1/wy), ceil(s2/wx)}
-static std::array<Value, 3> calculateDispatchWorkgroupCountFromResultShape(
-    Location loc, ArrayRef<Value> resultShape,
-    const std::array<Value, 3> &workgroupSize, OpBuilder &builder) {
-  if (resultShape.size() > 3) resultShape = resultShape.take_front(3);
-  SmallVector<Value, 4> reverseResultSize(reverse(resultShape));
-  Value one = builder.create<ConstantOp>(loc, builder.getIndexAttr(1));
-  reverseResultSize.resize(3, one);
-  return {
-      computeCeilDiv(loc, one, reverseResultSize[0], workgroupSize[0], builder),
-      computeCeilDiv(loc, one, reverseResultSize[1], workgroupSize[1], builder),
-      computeCeilDiv(loc, one, reverseResultSize[2], workgroupSize[2],
-                     builder)};
-}
-
-/// Calculates the number of workgroups to use based on the linearized shape of
-/// the result of the dispatch region. The `workgroupSize` is assumed to be of
-/// the form {wx, 1, 1}.  If the `resultShape` is {s0, s1, s2, ... sn}, then the
-/// number of workgroups is {ceil(s0*s1*s2*...*sn, wx)}
-static std::array<Value, 3>
-calculateDispatchWorkgroupCountFromLinearizedResultShape(
-    Location loc, ArrayRef<Value> resultShape,
-    const std::array<Value, 3> &workgroupSize, OpBuilder &builder) {
-  if (!mlir::matchPattern(workgroupSize[1], m_One()) ||
-      !mlir::matchPattern(workgroupSize[2], m_One())) {
-    emitError(loc,
-              "invalid workgroup size when computing workgroup count "
-              "based linearized result shape");
-    return {nullptr, nullptr, nullptr};
+/// The codegeneration emits a function `numWorkgroupsFn` for each entry point
+/// function. This function has arguments the !shapex.ranked_shape for all the
+/// input and output shaped types. Using this the function returns the number of
+/// workgroups to use. To use this function on the host side, generate the
+/// !shapex.ranked_shape values that describe the shape of the inputs and
+/// outputs of the dispatch region and "inline" the function body.
+static std::array<Value, 3> calculateWorkgroupCountFromNumWorkgroupsFn(
+    Location loc, FuncOp numWorkgroupsFn, IREE::HAL::InterfaceOp interface,
+    ArrayRef<Optional<TensorRewriteAdaptor>> operands,
+    ArrayRef<Optional<TensorRewriteAdaptor>> results, OpBuilder &builder) {
+  std::array<Value, 3> returnValue = {nullptr, nullptr, nullptr};
+  // TODO: This is really just inlining a function. For now assume that the
+  // `numWorkgroupsFn` has a single block to make inlining easier.
+  if (!numWorkgroupsFn || !llvm::hasSingleElement(numWorkgroupsFn))
+    return returnValue;
+  SmallVector<SmallVector<Value, 4>, 4> shapeValues;
+  shapeValues.reserve(operands.size() + results.size());
+  auto getShapeValuesFn =
+      [&](ArrayRef<Optional<TensorRewriteAdaptor>> values) -> LogicalResult {
+    for (auto val : values) {
+      if (!val) continue;
+      Optional<SmallVector<Value, 4>> shape = val->getShapeDims(builder);
+      if (!shape) return emitError(loc, "shape computation for operand failed");
+      shapeValues.push_back(shape.getValue());
+    }
+    return success();
+  };
+  if (failed(getShapeValuesFn(operands)) || failed(getShapeValuesFn(results)))
+    return returnValue;
+  BlockAndValueMapping mapper;
+  for (Operation &op : numWorkgroupsFn.front()) {
+    if (isa<mlir::ReturnOp>(op)) {
+      for (unsigned i = 0, e = std::min<unsigned>(3, op.getNumOperands());
+           i != e; ++i) {
+        returnValue[i] = mapper.lookupOrNull(op.getOperand(i));
+      }
+      break;
+    }
+    if (auto shapeOp = dyn_cast<Shape::RankedDimOp>(op)) {
+      if (BlockArgument arg = shapeOp.shape().dyn_cast<BlockArgument>()) {
+        auto &dimValues = shapeValues[arg.getArgNumber()];
+        mapper.map(arg, dimValues[shapeOp.getIndex()]);
+        continue;
+      }
+      return returnValue;
+    }
+    // If all its operands are mapped, clone it.
+    if (llvm::all_of(op.getOperands(), [&mapper](Value operand) {
+          return mapper.contains(operand);
+        })) {
+      builder.clone(op, mapper);
+      continue;
+    }
   }
-  Value one = builder.create<ConstantOp>(loc, builder.getIndexAttr(1));
-  Value linearizedSize = one;
-  for (Value dim : resultShape)
-    linearizedSize = builder.create<MulIOp>(loc, linearizedSize, dim);
-  return {computeCeilDiv(loc, one, linearizedSize, workgroupSize[0], builder),
-          one, one};
-}
-
-/// Calculates the number of workgroups to use for a dispatch region based on
-/// the value of `workgroupCountMethodAttr`. This is obtained from an attribute
-/// specified on the entry point functions that is added while lowering to
-/// SPIR-V.
-// TODO(ravishankarm): This method of using enums to specify methodology to
-// compute workgroup count is very hard to maintain. The best approach would be
-// that the lowering generates a function that is "inlined" here. Need to figure
-// out the signature of that function so that it covers all use cases.
-static std::array<Value, 3> calculateSPIRVDispatchWorkgroupCount(
-    Location loc, ArrayRef<Value> resultShape,
-    IntegerAttr workgroupCountMethodAttr,
-    const std::array<Value, 3> &workgroupSize, OpBuilder &builder) {
-  WorkgroupCountMethodology workgroupCountMethod =
-      static_cast<WorkgroupCountMethodology>(
-          workgroupCountMethodAttr.getValue().getZExtValue());
-  switch (workgroupCountMethod) {
-    case WorkgroupCountMethodology::Default:
-      return {nullptr, nullptr, nullptr};
-    case WorkgroupCountMethodology::LinearizeResultShape:
-      return calculateDispatchWorkgroupCountFromLinearizedResultShape(
-          loc, resultShape, workgroupSize, builder);
-    case WorkgroupCountMethodology::ResultShape:
-      return calculateDispatchWorkgroupCountFromResultShape(
-          loc, resultShape, workgroupSize, builder);
-  }
-  return {nullptr, nullptr, nullptr};
-}
-
-/// Gets the shape of the result from the dispatchState.
-static Optional<SmallVector<Value, 4>> getFirstResultShape(
-    Location loc, TargetBackend::DispatchState dispatchState,
-    OpBuilder &builder) {
-  if (dispatchState.results.empty()) return llvm::None;
-  Optional<TensorRewriteAdaptor> result = dispatchState.results[0];
-  SmallVector<Value, 4> resultShape;
-  // If the output is not a shaped type, assume it is a scalar, and return {1}.
-  if (!result) {
-    resultShape.push_back(
-        builder.create<ConstantOp>(loc, builder.getIndexAttr(1)));
-    return resultShape;
-  }
-
-  // TODO(ravishankarm): Using the result shape to get workgroup count, which
-  // involes using `getShapeDims,` results in the shape values being captured
-  // from outside of the switch statement in dynamic shape cases. This results
-  // in an error since switch statements cannot capture. For now, use the
-  // default path when the shape is dynamic.
-  if (!result->getTensorType().hasStaticShape()) return llvm::None;
-
-  return result->getShapeDims(builder);
+  return returnValue;
 }
 
 class VulkanSPIRVTargetBackend : public TargetBackend {
@@ -346,32 +304,18 @@
       auto workgroupSize = calculateDispatchWorkgroupSize(
           loc, spvModuleOp, spvFuncOp.sym_name(), workload, builder);
 
-      StringRef workgroupCountAttrName = getWorkgroupCountAttrName();
-      IntegerAttr workgroupCountAttr =
-          spvFuncOp.getAttrOfType<IntegerAttr>(workgroupCountAttrName);
-      if (!workgroupCountAttr)
-        return spvFuncOp.emitError("missing attribute ")
-               << workgroupCountAttrName;
-
-      // Assuming here that the shape of the first result value of the dispatch
-      // region is enough to calculate the number of workgroups. Either
-      // - All results have the same shape and the `workgroupCountMethod` is set
-      //   to WorkgroupCountMethodology::ResultShape, or
-      // - All the results have the same linearized shape and the
-      //   `workgourpCountMethod` is set to
-      //   WorkgroupCountMethodology::LinearizedResultShape.
-      Optional<SmallVector<Value, 4>> resultShape =
-          getFirstResultShape(loc, dispatchState, builder);
-
-      WorkgroupCountMethodology workgroupCountMethod =
-          static_cast<WorkgroupCountMethodology>(
-              workgroupCountAttr.getValue().getZExtValue());
+      FlatSymbolRefAttr numWorkgroupsFnAttr =
+          spvFuncOp.getAttrOfType<FlatSymbolRefAttr>(
+              getNumWorkgroupsFnAttrName());
 
       std::array<Value, 3> workgroupCount = {nullptr, nullptr, nullptr};
-      if (resultShape &&
-          workgroupCountMethod != WorkgroupCountMethodology::Default) {
-        workgroupCount = calculateSPIRVDispatchWorkgroupCount(
-            loc, *resultShape, workgroupCountAttr, workgroupSize, builder);
+      if (numWorkgroupsFnAttr) {
+        FuncOp numWorkgroupsFn = dyn_cast<FuncOp>(SymbolTable::lookupSymbolIn(
+            spvFuncOp.getParentOfType<ModuleOp>(), numWorkgroupsFnAttr));
+        if (!numWorkgroupsFn) return failure();
+        workgroupCount = calculateWorkgroupCountFromNumWorkgroupsFn(
+            loc, numWorkgroupsFn, executableOp.getInterfaceOp(),
+            dispatchState.operands, dispatchState.results, builder);
       } else {
         workgroupCount = calculateDispatchWorkgroupCount(
             loc, workload, workgroupSize, builder);