[mlir][linalg] Support for using output values in TC definitions.

[hanchung]

Summary

This will allow us to define select(pred, in, out) for TC ops, which is useful
for pooling ops.

Depends On D97269

[hanchung]

Please hold off for reviewing this atm.

[hanchung]

update the type to size_t

[hanchung]

I think this is ready now, I sent out the chain.

[antiagainst]

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

1090

Some documentation of what this means would be helpful.

[hanchung]

rebase and add comments

[hanchung]

rename def test and remove unused attributes

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

 ods_def<MatmulOp>
 implements_interface<LinalgContractionOpInterface> :
 def matmul(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
-  C(m, n) = std_addf<k>(std_mulf(A(m, k), B(k, n)));
+  C(m, n) = std_addf<k>(C(m, n), std_mulf(A(m, k), B(k, n)));
 }
 
 ods_def<MatmulColumnMajorOp>
 implements_interface<LinalgContractionOpInterface> :
 def matmul_column_major(A: f32(K, M), B: f32(N, K)) -> (C: f32(N, M)) {
-  C(n, m) = std_addf<k>(std_mulf(A(k, m), B(n, k)));
+  C(n, m) = std_addf<k>(C(n, m), std_mulf(A(k, m), B(n, k)));
 }
 
 ods_def<MatmulI8I8I32Op>
 implements_interface<LinalgContractionOpInterface> :
 def matmul_i8_i8_i32(A: i8(M, K), B: i8(K, N)) -> (C: i32(M, N)) {
   // TODO: ideally something closer to
   //   C(m, n) += cast<i32>(A(m, k)) * cast<i32>(B(k, n))
-  C(m, n) = std_addi<k>(std_sexti32(std_muli(A(m, k), B(k, n))));
+  C(m, n) = std_addi<k>(C(m, n), std_sexti32(std_muli(A(m, k), B(k, n))));
 }
 
 ods_def<MatvecOp>
 implements_interface<LinalgContractionOpInterface> :
 def matvec(A: f32(M, N), y: f32(N)) -> (x: f32(M)) {
-  x(m) = std_addf<n>(std_mulf(A(m, n), y(n)));
+  x(m) = std_addf<n>(x(m), std_mulf(A(m, n), y(n)));
 }
 
 ods_def<VecmatOp>
 implements_interface<LinalgContractionOpInterface> :
 def vecmat(y: f32(M), A: f32(M, N)) -> (x: f32(N)) {
-  x(n) = std_addf<m>(std_mulf(y(m), A(m, n)));
+  x(n) = std_addf<m>(x(n), std_mulf(y(m), A(m, n)));
 }
 
 ods_def<DotOp>
 implements_interface<LinalgContractionOpInterface> :
 def dot(A: f32(M), B: f32(M)) -> (C: f32()) {
-  C() = std_addf<m>(std_mulf(A(m), B(m)));
+  C() = std_addf<m>(C(), std_mulf(A(m), B(m)));
 }
 
 ods_def<BatchMatmulOp>
 implements_interface<LinalgContractionOpInterface> :
 def batch_matmul(A: f32(Batch, M, K), B: f32(Batch, K, N)) -> (C: f32(Batch, M, N)) {
-  C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(b, k, n)));
+  C(b, m, n) = std_addf<k>(C(b, m, n), std_mulf(A(b, m, k), B(b, k, n)));
 }
 
 ods_def<ConvWOp>:
 def conv_1d(I: f32(W), K: f32(KW)) -> (O: f32(W)) {
-  O(w) = std_addf<kw>(std_mulf(I(w + kw), K(kw)));
+  O(w) = std_addf<kw>(O(w), std_mulf(I(w + kw), K(kw)));
 }
 
 ods_def<ConvNWCOp>:
 def conv_1d_nwc(I: f32(N, W, C), K: f32(F, KW, C)) -> (O: f32(N, W, F)) {
-  O(n, w, f) = std_addf<kw>(std_mulf(I(n, w + kw, c), K(f, kw, c)));
+  O(n, w, f) = std_addf<kw>(O(n, w, f), std_mulf(I(n, w + kw, c), K(f, kw, c)));
 }
 
 ods_def<ConvNCWOp>:
 def conv_1d_ncw(I: f32(N, C, W), K: f32(F, C, KW)) -> (O: f32(N, F, W)) {
-  O(n, f, w) = std_addf<kw>(std_mulf(I(n, c, w + kw), K(f, c, kw)));
+  O(n, f, w) = std_addf<kw>(O(n, f, w), std_mulf(I(n, c, w + kw), K(f, c, kw)));
 }
 
 ods_def<ConvHWOp>:
 def conv_2d(I: f32(H, W), K: f32(KH, KW)) -> (O: f32(H, W)) {
-  O(h, w) = std_addf<kh, kw>(std_mulf(I(h + kh, w + kw), K(kh, kw)));
+  O(h, w) = std_addf<kh, kw>(O(h, w), std_mulf(I(h + kh, w + kw), K(kh, kw)));
 }
 
 ods_def<ConvNHWCOp>:
 def conv_2d_nhwc(I: f32(N, H, W, C), K: f32(F, KH, KW, C)) -> (O: f32(N, H, W, F)) {
-  O(n, h, w, f) = std_addf<kh, kw>(std_mulf(
-    I(n, h + kh, w + kw, c), K(f, kh, kw, c)));
+  O(n, h, w, f) = std_addf<kh, kw>(
+      O(n, h, w, f), std_mulf(I(n, h + kh, w + kw, c), K(f, kh, kw, c)));
 }
 
 ods_def<ConvNCHWOp>:
 def conv_2d_nchw(I: f32(N, C, H, W), K: f32(F, C, KH, KW)) -> (O: f32(N, F, H, W)) {
-  O(n, f, h, w) = std_addf<kh, kw>(std_mulf(
-    I(n, c, h + kh, w + kw), K(f, c, kh, kw)));
+  O(n, f, h, w) = std_addf<kh, kw>(
+      O(n, f, h, w), std_mulf(I(n, c, h + kh, w + kw), K(f, c, kh, kw)));
 }
 
 ods_def<ConvDHWOp>:
 def conv_3d(I: f32(D, H, W), K: f32(KD, KH, KW)) -> (O: f32(D, H, W)) {
-  O(d, h, w) = std_addf<kd, kh, kw>(std_mulf(
-    I(d + kd, h + kh, w + kw), K(kd, kh, kw)));
+  O(d, h, w) = std_addf<kd, kh, kw>(
+      O(d, h, w), std_mulf(I(d + kd, h + kh, w + kw), K(kd, kh, kw)));
 }
 
 ods_def<ConvNDHWCOp>:
 def conv_3d_ndhwc(I: f32(N, D, H, W, C), K: f32(F, KD, KH, KW, C)) -> (O: f32(N, D, H, W, F)) {
-  O(n, d, h, w, f) = std_addf<kd, kh, kw>(std_mulf(
-    I(n, d + kd, h + kh, w + kw, c), K(f, kd, kh, kw, c)));
+  O(n, d, h, w, f) = std_addf<kd, kh, kw>(
+      O(n, d, h, w, f),
+      std_mulf(I(n, d + kd, h + kh, w + kw, c), K(f, kd, kh, kw, c)));
 }
 
 ods_def<ConvNCDHWOp>:
 def conv_3d_ncdhw(I: f32(N, C, D, H, W), K: f32(F, C, KD, KH, KW)) -> (O: f32(N, F, D, H, W)) {
-  O(n, f, d, h, w) = std_addf<kd, kh, kw>(std_mulf(
-    I(n, c, d + kd, h + kh, w + kw), K(f, c, kd, kh, kw)));
+  O(n, f, d, h, w) = std_addf<kd, kh, kw>(
+      O(n, f, d, h, w),
+      std_mulf(I(n, c, d + kd, h + kh, w + kw), K(f, c, kd, kh, kw)));
 }
 
 ods_def<DepthwiseConvInputNHWCFilterHWCOp>:
 def depthwise_conv_2d_input_nhwc_filter_hwc
       (I: f32(N, IH, IW, C), K: f32(KH, KW, C))
    -> (O: f32(N, OH, OW, C))
   attr(strides: 2xi64)
 """A depth-wise 2-D convolution operation.
* `strides` is a 2-element vector attribute for window strides along the
   height/width dimension.
 
 The indexing maps for these three tensors contain 6 dimensions, following the
 order of (`N`, `OH`, `OW`, `C`, `KH`, `KW`).
 
 Note: this op only supports channel multiplier == 1.
 """
 {
-  O(n, oh, ow, c) = std_addf<kh, kw>(std_mulf(
-    I(n, oh * strides[0] + kh, ow * strides[1] + kw, c), K(kh, kw, c)));
+  O(n, oh, ow, c) = std_addf<kh, kw>(
+      O(n, oh, ow, c),
+      std_mulf(I(n, oh * strides[0] + kh, ow * strides[1] + kw, c),
+               K(kh, kw, c)));
 }
 
 ods_def<ConvInputNWCFilterWCFOp>:
 def conv_1d_input_nwc_filter_wcf(I: f32(N, W, C), K: f32(KW, C, F)) -> (O: f32(N, W, F))
   attr(strides: 1xi64, dilations: 1xi64)
 """ A 1-D convolution given NWC layout input and WCF layout filter.
 
 Computes a 1-D convolution given 3-D input and filter. The data layout
 of input is NWC and the data layout of filter is WCF.
 
 The indexing maps for these three tensors contain 5 dimensions, following the
 order of (`N`, `W`, `F`, `KW`, `C`).
 """
 {
   O(n, w, f) = std_addf<kw>(
+      O(n, w, f),
       std_mulf(I(n, w * strides[0] + kw * dilations[0], c), K(kw, c, f)));
 }
 
 ods_def<ConvInputNCWFilterWCFOp>:
 def conv_1d_input_ncw_filter_wcf(I: f32(N, C, W), K: f32(KW, C, F)) -> (O: f32(N, F, W))
   attr(strides: 1xi64, dilations: 1xi64)
 """ A 1-D convolution given NCW layout input and WCF layout filter.
 
 Computes a 1-D convolution given 3-D input and filter. The data layout
 of input is NCW and the data layout of filter is WCF.
 
 The indexing maps for these three tensors contain 5 dimensions, following the
 order of (`N`, `F`, `W`, `KW`, `C`).
 """
 {
   O(n, f, w) = std_addf<kw>(
+      O(n, f, w),
       std_mulf(I(n, c, w * strides[0] + kw * dilations[0]), K(kw, c, f)));
 }
 
 ods_def<ConvInputNHWCFilterHWCFOp>:
 def conv_2d_input_nhwc_filter_hwcf(I: f32(N, H, W, C), K: f32(KH, KW, C, F)) -> (O: f32(N, H, W, F))
   attr(strides: 2xi64, dilations: 2xi64)
 """ A 2-D convolution given NHWC layout input and HWCF layout filter.
 
 Computes a 2-D convolution given 4-D input and filter. The data layout
 of input is NHWC and the data layout of filter is HWCF.
 
 The indexing maps for these three tensors contain 7 dimensions, following the
 order of (`N`, `H`, `W`, `F`, `KH`, `KW`, `C`).
 """
 {
-  O(n, h, w, f) =
-      std_addf<kh, kw>(std_mulf(I(n, h * strides[0] + kh * dilations[0],
-                                  w * strides[1] + kw * dilations[1], c),
-                                K(kh, kw, c, f)));
+  O(n, h, w, f) = std_addf<kh, kw>(
+      O(n, h, w, f), std_mulf(I(n, h * strides[0] + kh * dilations[0],
 }
 
 ods_def<ConvInputNCHWFilterHWCFOp>:
 def conv_2d_input_nchw_filter_hwcf
     (I: f32(N, C, H, W), K: f32(KH, KW, C, F))
   -> (O: f32(N, F, H, W))
   attr(strides: 2xi64, dilations: 2xi64)
 """ A 2-D convolution given NCHW layout input and HWCF layout filter.
 
 Computes a 2-D convolution given 4-D input and filter. The data layout
 of input is NCHW and the data layout of filter is HWCF.
 
 The indexing maps for these three tensors contain 7 dimensions, following the
 order of (`N`, `F`, `H`, `W`, `KH`, `KW`, `C`).
 """
 {
-  O(n, f, h, w) =
-      std_addf<kh, kw>(std_mulf(I(n, c, h * strides[0] + kh * dilations[0],
-                                  w * strides[1] + kw * dilations[1]),
-                                K(kh, kw, c, f)));
+  O(n, f, h, w) = std_addf<kh, kw>(
+      O(n, f, h, w), std_mulf(I(n, c, h * strides[0] + kh * dilations[0],
 }
 
 ods_def<ConvInputNDHWCFilterDHWCFOp>:
 def conv_3d_input_ndhwc_filter_dhwcf
     (I: f32(N, D, H, W, C), K: f32(KD, KH, KW, C, F))
   -> (O: f32(N, D, H, W, F))
   attr(strides: 3xi64, dilations: 3xi64)
 """ A 3-D convolution given NDHWC layout input and DHWCF layout filter.
 
 Computes a 3-D convolution given 5-D input and filter. The data layout
 of input is NDHWC and the data layout of filter is DHWCF.
 
 The indexing maps for these three tensors contain 9 dimensions, following the
 order of (`N`, `D`, `H`, `W`, `F`, `KD`, `KH`, `KW`, `C`).
 """
 {
-  O(n, d, h, w, f) =
-      std_addf<kd, kh, kw>(std_mulf(I(n, d * strides[0] + kd * dilations[0],
-                                      h * strides[1] + kh * dilations[1],
-                                      w * strides[2] + kw * dilations[2], c),
-                                    K(kd, kh, kw, c, f)));
+  O(n, d, h, w, f) = std_addf<kd, kh, kw>(
+      O(n, d, h, w, f), std_mulf(I(n, d * strides[0] + kd * dilations[0],
 }
 
 ods_def<ConvInputNCDHWFilterDHWCFOp>:
 def conv_3d_input_ncdhw_filter_dhwcf
     (I: f32(N, C, D, H, W), K: f32(KD, KH, KW, C, F))
   -> (O: f32(N, F, D, H, W))
   attr(strides: 3xi64, dilations: 3xi64)
 """ A 3-D convolution given NCDHW layout input and DHWCF layout filter.
 
 Computes a 3-D convolution given 5-D input and filter. The data layout
 of input is NCDHW and the data layout of filter is DHWCF.
 
 The indexing maps for these three tensors contain 9 dimensions, following the
 order of (`N`, `F`, `D`, `H`, `W`, `KD`, `KH`, `KW`, `C`).
 """
 {
-  O(n, f, d, h, w) = std_addf<kd, kh, kw>(std_mulf(
-      I(n, c, d * strides[0] + kd * dilations[0],
-        h * strides[1] + kh * dilations[1], w * strides[2] + kw * dilations[2]),
+  O(n, f, d, h, w) = std_addf<kd, kh, kw>(
+      O(n, f, d, h, w), std_mulf(I(n, c, d * strides[0] + kd * dilations[0],
+                                   h * strides[1] + kh * dilations[1],
+                                   w * strides[2] + kw * dilations[2]),
                                  K(kd, kh, kw, c, f)));
 }
+}

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc

 //       IMPL:  void Test1Op::regionBuilder(Block &block, ValueRange captures) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 ods_def<Test1Op> :
 def test1(A: f32(M, K), B: f32(K)) -> (C: f32(M)) {
-  C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
+  C(m) = std_addf<k>(C(m), std_mulf(A(m, k), B(k)));
 }
 
 // ODS-LABEL: def Test2Op : LinalgStructuredBase_Op<"test2", [
 //  ODS-NEXT:   AttrSizedOperandSegments
 //  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
 // IMPL-LABEL:  ArrayAttr Test2Op::iterator_types() {
 //       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
 //       IMPL:  ArrayAttr Test2Op::indexing_maps() {
 //       IMPL:  AffineMap::get(3, 3, {d0, d2}, context)
 //       IMPL:  AffineMap::get(3, 3, {d2, d1}, context)
 //       IMPL:  AffineMap::get(3, 3, {d0, d1}, context)
 //
 //       IMPL:  Test2Op::regionBuilder(Block &block, ValueRange captures) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 ods_def<Test2Op> :
 def test2(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
-  C(m, n) = std_addf<k>(std_mulf(A(m, k), B(k, n)));
+  C(m, n) = std_addf<k>(C(m, n), std_mulf(A(m, k), B(k, n)));
 }
 
 // ODS-LABEL: def Test3Op : LinalgStructuredBase_Op<"test3", [
 //  ODS-NEXT:   AttrSizedOperandSegments
 //  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
 // IMPL-LABEL:  ArrayAttr Test3Op::iterator_types() {
 //       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
 //       IMPL:  ArrayAttr Test3Op::indexing_maps() {
 //       IMPL:  AffineMap::get(4, 4, {d0, d1, d3}, context)
 //       IMPL:  AffineMap::get(4, 4, {d3, d2}, context)
 //       IMPL:  AffineMap::get(4, 4, {d0, d1, d2}, context)
 //
 //       IMPL:  Test3Op::regionBuilder(Block &block, ValueRange captures) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 ods_def<Test3Op> :
 def test3(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
-  C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
+  C(b, m, n) = std_addf<k>(C(b, m, n), std_mulf(A(b, m, k), B(k, n)));
 }
 
 // Test attribute definitions
 // ODS-LABEL: def Test4Op
 // ODS: F32ArrayAttr:$array_attr,
 // ODS: F32:$f32_attr,
 // ODS: RankedF32ElementsAttr<[4]>:$fvec_attr,
 // ODS: I32:$i32_attr,
attr(
   f32_attr: f32,
   i32_attr: i32,
   i64_attr: i64,
   fvec_attr: 4xf32,
   ivec_attr: 5x6xi32,
   array_attr : f32[],
   optional_attr? : f32
 ) {
-  C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
+  C(b, m, n) = std_addf<k>(C(b, m, n), std_mulf(A(b, m, k), B(k, n)));
 }
 
 // Test attribute usage in affine expressions
 // IMPL-LABEL: ArrayAttr Test5Op::indexing_maps() {
 // IMPL: auto cst0 = getAffineConstantExpr(strides().getValue<int>({ 0 }), context);
 // IMPL: auto cst1 = getAffineConstantExpr(strides().getValue<int>({ 1 }), context);
 // IMPL: auto map0 = AffineMap::get(7, 9, {d0, d1 * s7 + d4, d2 * s8 + d5, d6}, context);
 // IMPL: map0 = map0.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);

 def test6(A: f32(M, K), B: f32(K)) -> (C: f32(M))
 """
 My magic op.
 
 It has two inputs.
 It has one output.
 """
 {
-  C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
+  C(m) = std_addf<k>(C(m), std_mulf(A(m, k), B(k)));
 }
 
 // Test attribute builder
 // ODS-LABEL: def Test7Op
 // ODS:         OpBuilderDAG<
 // ODS:           (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
 // ODS:            "ValueRange":$outputs, "Attribute":$attr_a, "Attribute":$attr_b)
 // ODS:           $_state.addAttribute("attr_a", attr_a);
 // ODS:           $_state.addAttribute("attr_b", attr_b);
 //
 ods_def<Test7Op>:
 def test7(A: f32(M, K), B: f32(K)) -> (C: f32(M))
      attr(attr_a: f32, attr_b: 4xi32)
 {
-  C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
+  C(m) = std_addf<k>(C(m), std_mulf(A(m, k), B(k)));
+}
+
+// Test output arg order.
+// IMPL-LABEL:  void Test8Op::regionBuilder(Block &block, ValueRange captures) {
+//       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
+//       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
+//       IMPL:  Value [[e:.*]] = std_subf([[d]], [[c]]);
+//       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
+ods_def<Test8Op>:
+def test7(A: f32(M, K), B: f32(K)) -> (C: f32(M))
+     attr(attr_a: f32, attr_b: 4xi32)
+{
+  C(m) = std_subf<k>(std_mulf(A(m, k), B(k)), C(m));
 }

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

public:
   /// Parse the information for a tensor def.
   /// All the affine-expr must be dimensionless (i.e. contain only expressions
   /// involving symbols and constants), but can otherwise contain arbitrary
   /// affine expressions.
   LogicalResult parseTensorDef(bool isOutput);
 
   /// Parses a tensor use.
   struct ComprehensionParsingState {
+    size_t numArgs;
     AffineDimList dims;
     SmallVector<std::unique_ptr<Expression>, 4> expressions;
     llvm::DenseMap<TensorUse, unsigned> orderedTensorArgs;
   };
   LogicalResult parseTensorUse(TensorUse &result,
                                ComprehensionParsingState &state);
 
   /// Parses an attribute definition.
LogicalResult TCParser::parseExpression(TensorUse currentDefinition,
   // Check if it has a reduction set, discover dimensions eagerly.
   if (parser.curToken.is(Token::Kind::lt)) {
     auto iters = parseAffineExprs(EagerDiscoveryMode::Dimensions, state.dims,
                                   Token::Kind::lt, Token::Kind::gt);
     for (auto iter : iters)
       reductionDims.push_back(iter.cast<AffineDimExpr>().getPosition());
   }
 
-  // If this op is a reduction, it's first argument is the `currentDefinition`
-  // tensor use.
-  if (!reductionDims.empty())
-    expressions.push_back(std::make_unique<TensorUse>(currentDefinition));
-  LLVM_DEBUG(llvm::dbgs() << "op: " << opOrTensor << "\n");
 
   auto parseExpr = [&]() -> LogicalResult {
     std::unique_ptr<Expression> e;
     if (failed(parseExpression(currentDefinition, e, state)))
       return failure();
     expressions.push_back(std::move(e));
     return success();
   };
visitPostorder(*pExpr, [&](const Expression &e) {
       if (failed || !pUse)
         return;
       auto &use = *pUse;
       LLVM_DEBUG(llvm::dbgs()
                  << "\nuse: " << use.tensorId << " map: " << use.indexingMap);
       auto tensorIter = registeredTensors.find(use.tensorId);
       assert(tensorIter != registeredTensors.end() && "unregistered tensor");
       auto &tensor = tensorIter->getValue();
-      if (tensor.indexingMap && state.orderedTensorArgs.count(use) == 0) {
+      if (tensor.indexingMap && state.orderedTensorArgs.count(use) == 0 &&
+          tensor.indexingMap.getResults() != use.indexingMap.getResults()) {
         LLVM_DEBUG(llvm::dbgs() << "\nexisting: " << tensor.indexingMap);
         (void)parser.emitError(
             "Unexpected multi-read of a tensor with different accesses");
         failed = true;
         return;
       }
       seenDefs.insert(use.tensorId);
       tensor.indexingMap = use.indexingMap;
       state.orderedTensorArgs[use] = tensor.index;
     });
+  state.numArgs = seenDefs.size();
   if (failed)
     return failure();
 
   return success();
 }
 
 /// Parse and print the information for a ODS def.
 ///
if (!registeredAttrs.empty()) {
     attrMethods = R"(
       bool hasDynamicIndexingMaps();
       LogicalResult verifyIndexingMapRequiredAttributes();
     )";
   }
 
   // Finally put everything together.
   os << llvm::formatv(header, cppOpName, linalgOpName, interfaceNameList, doc,
-                      attrList, state.orderedTensorArgs.size(), attrBuilder,
-                      attrMethods);
+                      attrList, state.numArgs, attrBuilder, attrMethods);
 }
 
 /// Print the C++ StructuredOpsInterface impl of `iterator_types`.
 void TCParser::printReferenceIterators(llvm::raw_ostream &os,
                                        StringRef cppOpName,
                                        ComprehensionParsingState &state) {
   const char *referenceReferenceIteratorsFmt =
       R"FMT(
void TCParser::printReferenceIndexingMaps(llvm::raw_ostream &os,
   ss.flush();
 
   // 3. Get the list of affine maps for each input/output. The AffineExpr use
   // the common arithmetic operators on AffineExpr. These affine maps will
   // replace the `{2}` placeholder.
   std::string mapsStr;
   llvm::raw_string_ostream mapsStringStream(mapsStr);
 
-  SmallVector<TensorUse, 4> orderedUses(state.orderedTensorArgs.size());
+  SmallVector<TensorUse, 4> orderedUses(state.numArgs);
   for (const auto &it : state.orderedTensorArgs)
     orderedUses[it.second] = it.first;
 
   // Create a list of all symbols.
   SmallVector<std::string, 4> symbolReplacements;
   symbolReplacements.reserve(symbols.size());
   for (unsigned i = 0; i < symbols.size(); ++i) {
     const char *symFmt =
void TCParser::printReferenceIndexingMaps(llvm::raw_ostream &os,
   // 4. Apply format to 1. using 2. and 3.
   os << llvm::formatv(referenceIndexingMapsFmt, cppOpName, dimsStr, mapsStr,
                       llvm::join(mapList, ", "));
 }
 
 /// Print the C++ StructuredOpsInterface impl of `regionBuilder`.
 void TCParser::printRegionBuilder(llvm::raw_ostream &os, StringRef cppOpName,
                                   ComprehensionParsingState &state) {
-  unsigned count = state.orderedTensorArgs.size();
+  unsigned count = state.numArgs;
   llvm::DenseMap<const TensorExpr *, unsigned> subExprsMap;
   std::function<void(llvm::raw_ostream & os, const Expression &)> printExpr;
   printExpr = [&](llvm::raw_ostream &os, const Expression &e) -> void {
     if (auto *pUse = dyn_cast<TensorUse>(&e)) {
       os << "_" << state.orderedTensorArgs.find(*pUse)->second;
       return;
     }
     auto *pTensorExpr = cast<TensorExpr>(&e);
void {0}::regionBuilder(Block &block, ValueRange captures) {
     {2}
     (linalg_yield(ValueRange{ {3} }));
   })FMT";
 
   unsigned idx = 0;
   std::string valueHandleStr;
   llvm::raw_string_ostream valueHandleStringStream(valueHandleStr);
   llvm::interleaveComma(
-      state.orderedTensorArgs, valueHandleStringStream, [&](auto) {
+      llvm::seq<int>(0, state.numArgs), valueHandleStringStream, [&](auto) {
         valueHandleStringStream << "_" << idx << "(args[" << idx << "])";
         idx++;
       });
 
   std::string expressionsStr;
   llvm::raw_string_ostream expressionStringStream(expressionsStr);
   for (auto &expr : state.expressions)
     visitPostorder(*expr, [&](const Expression &e) {