torch-mlir/include/torch-mlir-dialects/Dialect/TMTensor/IR/ScalarLoopOpInterface.td

//===-------------------------------------------------------*- tablegen -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// Also available under a BSD-style license. See LICENSE.
//
//===----------------------------------------------------------------------===//

#ifndef TORCH_MLIR_DIALECT_TMTENSOR_SCALARLOOPOPINTERFACE
#define TORCH_MLIR_DIALECT_TMTENSOR_SCALARLOOPOPINTERFACE

include "mlir/IR/OpBase.td"

def ScalarLoopOpInterface : OpInterface<"ScalarLoopOpInterface"> {
  let description = [{
    Interface for allowing operations to expose information needed to
    lower it to for loops
  }];
  let cppNamespace = "::mlir::torch::TMTensor";
  let methods = [
      InterfaceMethod<
        /*desc=*/[{
          Returns the destination operands. For op with `memref`
          operands, this is the result buffers. For op with `tensor`
          operands, this is the operands that contain the initial
          value for the result. These are "tied" to the result
          buffers. For example, for a `LinalgOp`/`TMTensor` ops, it
          is the `outs` parameters. For `tensor.insert_slice`, it is
          the `dest` parameter.
        }],
        /*retType=*/"SmallVector<Value>",
        /*methodName=*/"getDestinationOperands",
        /*args=*/(ins "OpBuilder &":$b),
        /*methodBody=*/"",
        /*defaultImplementation=*/"return ValueRange{};"
      >,
      InterfaceMethod<
        /*desc=*/[{
          Returns a list of `IteratorType`s that describe the number of
          loops and the iterator types of the operation. The list is
          expected to use
          `utils::IteratorType::parallel`/`utils::IteratorType::reduction`
          from MLIR Structured Op Utils.
        }],
        /*retType=*/"SmallVector<utils::IteratorType>",
        /*methodName=*/"getLoopIteratorTypes"
      >,
      InterfaceMethod<
        /*desc=*/[{
          Returns a list of ranges that describe the loop bounds and
          step for the loops of the operation.
        }],
        /*retTy=*/"SmallVector<Range>",
        /*methodName=*/"getIterationDomain",
        /*args=*/(ins "OpBuilder &":$b)
      >,
      InterfaceMethod<
        /*desc=*/[{
          Generates the loop body implementation. Assume that all the parallel
          loops and reduction loops are created and the insertion point of the
          build is set to the innermost of the loop. This method implements the
          loop body IRs.
        }],
        /*retType=*/"LogicalResult",
        /*methodName=*/"generateScalarImplementation",
        /*args=*/(ins
            "OpBuilder &":$b,
            "Location ":$loc,
            "ValueRange ":$ivs),
        /*methodBody=*/"",
        /*defaultImplementation=*/[{
          return failure();
        }]
      >
  ];
}

#endif  // TORCH_MLIR_DIALECT_TMTENSOR_SCALARLOOPOPINTERFACES